Hey Team,
Seeking advice on code efficiency.
I am attempting to solve a hyperelastic mixed-element problem over different displacement quantities. However, even with this very simple MWE the code is quite slow.
MWE:
# ∆ Dolfin
import numpy as np
import ufl
from mpi4py import MPI
from basix.ufl import element, mixed_element
from dolfinx import log, io, default_scalar_type, mesh
from dolfinx.fem import Function, functionspace, dirichletbc, locate_dofs_topological
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver
from ufl import (Identity, grad, TestFunctions, split, det, as_tensor, variable, dx)
# ∆ Parameters
DIM = 3
ORDER = 2
TOL = 1e-5
QUAD_DEGREE = 4
# ∆ Mesh
domain = mesh.create_unit_cube(MPI.COMM_WORLD, 10, 10, 10)
domain.name = "mesh"
# ∆ Function space
P2 = element("Lagrange", domain.basix_cell(), ORDER, shape=(domain.geometry.dim,))
P1 = element("Lagrange", domain.basix_cell(), ORDER-1)
Mxs = functionspace(mesh=domain, element=mixed_element([P2, P1]))
Tes = functionspace(mesh=domain, element=("Lagrange", ORDER, (DIM, DIM)))
# ∆ Define subdomains
V, _ = Mxs.sub(0).collapse()
# ∆ Trial & Test
u_p = Function(Mxs)
v, q = TestFunctions(Mxs)
u, p = split(u_p)
# ∆ Kinematics
I = Identity(DIM)
F = variable(I + grad(u))
C = variable(F.T * F)
E = as_tensor(0.5 * (C - I))
J = det(F)
# ∆ Kinematics
I = ufl.variable(ufl.Identity(DIM))
F = ufl.variable(I + ufl.grad(u))
C = ufl.variable(F.T * F)
Ic = ufl.variable(ufl.tr(C))
IIc = ufl.variable((Ic**2 - ufl.inner(C,C))/2)
J = ufl.variable(ufl.det(F))
psi = 1 * (Ic - 3) + 0 *(IIc - 3)
gamma1 = ufl.diff(psi, Ic) + Ic * ufl.diff(psi, IIc)
gamma2 = -ufl.diff(psi, IIc)
firstPK = 2 * F * (gamma1*I + gamma2*C) + p * J * ufl.inv(F).T
cau = (1/J * firstPK * F).T
# ∆ Residual
dx = ufl.Measure(integral_type="dx", domain=domain, metadata={"quadrature_degree": QUAD_DEGREE})
R = ufl.inner(ufl.grad(v), firstPK) * dx + q * (J - 1) * dx
# ∆ Solver
problem = NonlinearProblem(R, u_p, [])
solver = NewtonSolver(MPI.COMM_WORLD, problem)
solver.atol = TOL
solver.rtol = TOL
solver.max_it = 50
solver.convergence_criterion = "incremental"
# ∆ Output
displacement = Function(V)
displacement.name = "Displacement"
xdmf = io.VTXWriter(domain.comm, "disp.bp", displacement, engine="BP4")
X, Y, Z = 0, 1, 2
def left(x):
return np.isclose(x[0], 0)
def right(x):
return np.isclose(x[0], max(x[0]))
fdim = domain.topology.dim - 1
left_facets = mesh.locate_entities_boundary(domain, fdim, left)
right_facets = mesh.locate_entities_boundary(domain, fdim, right)
marked_facets = np.hstack([left_facets, right_facets])
marked_values = np.hstack([np.full_like(left_facets, 1), np.full_like(right_facets, 2)])
sorted_facets = np.argsort(marked_facets)
facet_tag = mesh.meshtags(domain, fdim, marked_facets[sorted_facets], marked_values[sorted_facets])
# ∆ Setup boundary terms
xx0 = locate_dofs_topological(Mxs.sub(0).sub(X), domain.topology.dim - 1, facet_tag.find(1))
xx1 = locate_dofs_topological(Mxs.sub(0).sub(X), domain.topology.dim - 1, facet_tag.find(2))
yx0 = locate_dofs_topological(Mxs.sub(0).sub(Y), domain.topology.dim - 1, facet_tag.find(1))
yx1 = locate_dofs_topological(Mxs.sub(0).sub(Y), domain.topology.dim - 1, facet_tag.find(2))
zx0 = locate_dofs_topological(Mxs.sub(0).sub(Z), domain.topology.dim - 1, facet_tag.find(1))
zx1 = locate_dofs_topological(Mxs.sub(0).sub(Z), domain.topology.dim - 1, facet_tag.find(2))
# ∆ Set boundaries
d_yy0 = dirichletbc(default_scalar_type(0.0), yx0, Mxs.sub(0).sub(Y))
d_yy1 = dirichletbc(default_scalar_type(0.0), yx1, Mxs.sub(0).sub(Y))
d_zy0 = dirichletbc(default_scalar_type(0.0), zx0, Mxs.sub(0).sub(Z))
d_zy1 = dirichletbc(default_scalar_type(0.0), zx1, Mxs.sub(0).sub(Z))
log.set_log_level(log.LogLevel.INFO)
# ∆ Iterate
for ii, kk in enumerate([0, 5, 10, 15, 20]):
# ∆ Apply displacements as boundary conditions
du = kk / 100
d_xy0 = dirichletbc(default_scalar_type(du), xx0, Mxs.sub(0).sub(X))
d_xy1 = dirichletbc(default_scalar_type(0.0), xx1, Mxs.sub(0).sub(X))
bc = [d_xy0, d_yy0, d_zy0, d_xy1, d_yy1, d_zy1]
problem.bcs = bc
# ∆ Solve
try:
n_iter, conv = solver.solve(u_p)
print(f"Solved {kk}%: iterations={n_iter}, residual={conv}")
except RuntimeError:
print(f"Failed to converge at {kk}%")
break
# ∆ Store output
displacement.interpolate(u_p.sub(0).collapse())
xdmf.write(ii)
xdmf.close()
I have access to quite a powerful HPC for the real problem but solve times have exceeded 24 hours. I assume there is some level of computational ignorance on my end.
Appreciate any advice 
Versions:
fenics-basix 0.9.0 py313hd7981f8_2 conda-forge
fenics-basix-nanobind-abi 0.2.1.18 h9b63b7c_2 conda-forge
fenics-dolfinx 0.9.0 py313h12399d7_114 conda-forge
fenics-ffcx 0.9.0 pyh2e48890_0 conda-forge
fenics-libbasix 0.9.0 hf7ae0cd_2 conda-forge
fenics-libdolfinx 0.9.0 py313hd0b34b2_114 conda-forge
fenics-ufcx 0.9.0 hb7f7608_0 conda-forge
fenics-ufl 2024.2.0 pyhd8ed1ab_1 conda-forge
A snippet of the more complicated variational form (non-MWE) is provided below if interested
# ∆ Load mesh data and set up function spaces
domain, _, ft = io.gmshio.read_from_msh(filename=file, comm=MPI.COMM_WORLD, rank=0, gdim=DIM)
P2 = element("Lagrange", domain.basix_cell(), ORDER, shape=(domain.geometry.dim,))
P1 = element("Lagrange", domain.basix_cell(), ORDER-1)
Mxs = functionspace(mesh=domain, element=mixed_element([P2, P1]))
Tes = functionspace(mesh=domain, element=("Lagrange", ORDER, (DIM, DIM)))
# ∆ Define subdomains
V, _ = Mxs.sub(0).collapse()
P, _ = Mxs.sub(1).collapse()
# ∆ Determine coordinates of space and create mapping tree
x_n = Function(V)
coords = np.array(x_n.function_space.tabulate_dof_coordinates()[:])
tree = KDTree(coords)
if t != "test":
# ∆ Setup functions for assignment
ori, z_data = Function(V), Function(V)
# ∆ Assign angle and z-disc data
azi, ele, zs = angle_assign(t, coords)
# ∆ Store angles
CA, CE = np.cos(azi), np.cos(ele)
SA, SE = np.sin(azi), np.sin(ele)
# ∆ Create interpolate functions
# µ Basis vector 1
def nu_1(phi_xyz):
_, idx = tree.query(phi_xyz.T, k=1)
return np.array([CA[idx]*CE[idx], SA[idx]*CE[idx], -SE[idx]])
# µ Basis vector 2
def nu_2(phi_xyz):
_, idx = tree.query(phi_xyz.T, k=1)
return np.array([-SA[idx], CA[idx], np.zeros_like(CA[idx])])
# µ Basis vector 3
def nu_3(phi_xyz):
_, idx = tree.query(phi_xyz.T, k=1)
return np.array([CA[idx]*SE[idx], SA[idx]*SE[idx], CE[idx]])
# ∆ Create z_disc id data
z_arr = z_data.x.array.reshape(-1, 3)
z_arr[:, 0], z_arr[:, 1], z_arr[:, 2] = zs, azi, ele
z_data.x.array[:] = z_arr.reshape(-1)
# ∆ Create angular orientation vector
ori.interpolate(nu_1)
# ∆ Save foundations data
z_data.name = "Node Mapping"
ori.name = "Orientation Vectors"
with io.VTXWriter(MPI.COMM_WORLD, f"_bp/_{t}/_ZSN.bp", z_data, engine="BP4") as fz:
fz.write(0)
fz.close()
with io.VTXWriter(MPI.COMM_WORLD, f"_bp/_{t}/_ORI.bp", ori, engine="BP4") as fo:
fo.write(0)
fo.close()
# ∆ Create push tensor function
Push = Function(Tes)
# ∆ Define push interpolation
def forward(phi_xyz):
_, idx = tree.query(phi_xyz.T, k=1)
f00, f01, f02 = CA[idx]*CE[idx], -SA[idx], CA[idx]*SE[idx]
f10, f11, f12 = SA[idx]*CE[idx], CA[idx], SA[idx]*SE[idx]
f20, f21, f22 = -SE[idx], np.zeros_like(CE[idx]), CE[idx]
return np.array([f00, f01, f02, f10, f11, f12, f20, f21, f22])
# ∆ Interpolate Push as Forward transform
Push.interpolate(forward)
else:
# ∆ Push as Forward transform
Push = ufl.Identity(DIM)
# ∆ Load key function variables
mx = Function(Mxs)
v, q = ufl.TestFunctions(Mxs)
u, p = ufl.split(mx)
u_nu = Push * u
# ∆ Kinematics Setup
i, j, k, l, a, b = ufl.indices(6)
I = ufl.Identity(DIM)
F = ufl.variable(I + ufl.grad(u_nu))
if t != "test":
# ∆ Metric tensors
# µ [UNDERFORMED] Covariant basis vectors
A1, A2, A3 = Function(V), Function(V), Function(V)
# ¬ create base 1
A1.interpolate(nu_1)
# ¬ create base 2
A2.interpolate(nu_2)
# ¬ create base 3
A3.interpolate(nu_3)
# µ [UNDERFORMED] Metric tensors
G_v = ufl.as_tensor([
[ufl.dot(A1, A1), ufl.dot(A1, A2), ufl.dot(A1, A3)],
[ufl.dot(A1, A2), ufl.dot(A2, A2), ufl.dot(A2, A3)],
[ufl.dot(A1, A3), ufl.dot(A2, A3), ufl.dot(A3, A3)]
])
G_v_inv = ufl.inv(G_v)
# µ [DEFORMED] Metric covariant tensors
g_v = ufl.as_tensor([
[ufl.dot(F * A1, F * A1), ufl.dot(F * A1, F * A2), ufl.dot(F * A1, F * A3)],
[ufl.dot(F * A2, F * A1), ufl.dot(F * A2, F * A2), ufl.dot(F * A2, F * A3)],
[ufl.dot(F * A3, F * A1), ufl.dot(F * A3, F * A2), ufl.dot(F * A3, F * A3)]
])
# ∆ Christoffel symbols
Gamma = ufl.as_tensor(
0.5 * G_v_inv[k, l] * (ufl.grad(G_v[j, l])[i] + ufl.grad(G_v[i, l])[j] - ufl.grad(G_v[i, j])[l]),
(i, j, k)
)
# ∆ Covariant derivative
covDev = ufl.as_tensor(ufl.grad(v)[i, j] + Gamma[i, k, j] * v[k], (i, j))
else:
# ∆ Covariant derivative
covDev = ufl.as_tensor(ufl.grad(v)[i, j], (i, j))
# ∆ Kinematics Tensors
C = ufl.variable(F.T * F)
B = ufl.variable(F * F.T)
if t != "test":
E = ufl.as_tensor(0.5 * (g_v - G_v))
else:
E = ufl.as_tensor(0.5 * (C - I))
J = ufl.det(F)
# ∆ Extract Constitutive terms
b0, bf, bt = gcc
# ∆ Exponent term
Q = (
bf * E[0,0]**2 + bt *
(
E[1,1]**2 + E[2,2]**2 + E[1,2]**2 + E[2,1]**2 +
E[0,1]**2 + E[1,0]**2 + E[0,2]**2 + E[2,0]**2
)
)
# ∆ Seond Piola-Kirchoff
if t != "test":
SPK = b0/4 * ufl.exp(Q) * ufl.as_matrix([
[4*bf*E[0,0], 2*bt*(E[1,0] + E[0,1]), 2*bt*(E[2,0] + E[0,2])],
[2*bt*(E[0,1] + E[1,0]), 4*bt*E[1,1], 2*bt*(E[2,1] + E[1,2])],
[2*bt*(E[0,2] + E[2,0]), 2*bt*(E[1,2] + E[2,1]), 4*bt*E[2,2]],
]) - p * G_v_inv
else:
SPK = b0/4 * ufl.exp(Q) * ufl.as_matrix([
[4*bf*E[0,0], 2*bt*(E[1,0] + E[0,1]), 2*bt*(E[2,0] + E[0,2])],
[2*bt*(E[0,1] + E[1,0]), 4*bt*E[1,1], 2*bt*(E[2,1] + E[1,2])],
[2*bt*(E[0,2] + E[2,0]), 2*bt*(E[1,2] + E[2,1]), 4*bt*E[2,2]],
]) - p * I
# ∆ Residual
dx = ufl.Measure(integral_type="dx", domain=domain, metadata={"quadrature_degree": QUADRATURE})
R = ufl.as_tensor(SPK[a, b] * F[j, b] * covDev[j, a]) * dx + q * (J - 1) * dx