Static condensation of linear elasticity
Copyright (C) 2020 Michal Habera and Andreas Zilian
Download sources
This demo solves a Cook’s plane stress elasticity test in a mixed space formulation. The test is a sloped cantilever under upward traction force at free end. Static condensation of internal (stress) degrees-of-freedom is demonstrated.
This demo illustrates how to:
Use static condensation with numba on variational forms made with UFL.
Extracting JIT compiled C-kernels using
ffcx_jit
This demo requires more modules than usual, as it uses numba for
efficient static condensation.
from pathlib import Path
from mpi4py import MPI
from petsc4py import PETSc
import cffi
import numba
import numba.core.typing.cffi_utils as cffi_support
import numpy as np
import ufl
from basix.ufl import element
from dolfinx import default_real_type, default_scalar_type, geometry
from dolfinx.fem import (
Form,
Function,
IntegralType,
dirichletbc,
form,
form_cpp_class,
functionspace,
locate_dofs_topological,
)
from dolfinx.fem.petsc import apply_lifting, assemble_matrix, assemble_vector
from dolfinx.io import XDMFFile
from dolfinx.jit import ffcx_jit
from dolfinx.mesh import locate_entities_boundary, meshtags
from ffcx.codegeneration.utils import empty_void_pointer
from ffcx.codegeneration.utils import numba_ufcx_kernel_signature as ufcx_signature
rtype = default_real_type
dtype = default_scalar_type
if np.issubdtype(rtype, np.float32): # type: ignore
print("float32 not yet supported for this demo.")
exit(0)
We start by reading in the Cook’s mesh cooks_tri_mesh.xdmf
using XDMFFile.read_mesh.
Note that the mesh is written in plain-text format, which means we use
XDMFFile.Encoding.ASCII.
infile = XDMFFile(
MPI.COMM_WORLD,
Path(Path(__file__).parent, "data", "cooks_tri_mesh.xdmf"),
"r",
encoding=XDMFFile.Encoding.ASCII,
)
msh = infile.read_mesh(name="Grid")
infile.close()
We create the Stress (Se) and displacement (Ue) elements and corresponding function spaces. Note that the stress element is symmetric.
gdim = msh.geometry.dim
Se = element("DG", msh.basix_cell(), 1, shape=(gdim, gdim), symmetry=True, dtype=rtype) # type: ignore
Ue = element("Lagrange", msh.basix_cell(), 2, shape=(gdim,), dtype=rtype) # type: ignore
S = functionspace(msh, Se)
U = functionspace(msh, Ue)
Next, we define the trial and test functions for stress and displacement,
sigma, tau = ufl.TrialFunction(S), ufl.TestFunction(S)
u, v = ufl.TrialFunction(U), ufl.TestFunction(U)
Locate all facets at the free end and assign them value 1. Sort the
facet indices (requirement for constructing MeshTags).
tdim = msh.topology.dim
free_end_facets = np.sort(locate_entities_boundary(msh, tdim - 1, lambda x: np.isclose(x[0], 48.0)))
mt = meshtags(msh, tdim - 1, free_end_facets, 1)
Next, we create an integration measure with the facet markers.
ds = ufl.Measure("ds", subdomain_data=mt)
Homogeneous boundary condition in displacement
Displacement BC is applied to
the left side
left_facets = locate_entities_boundary(msh, tdim - 1, lambda x: np.isclose(x[0], 0.0))
bdofs = locate_dofs_topological(U, tdim - 1, left_facets)
bc = dirichletbc(u_bc, bdofs)
Elastic stiffness tensor and Poisson ratio
E, nu = 1.0, 1.0 / 3.0
def sigma_u(u):
"""Constitutive relation for stress-strain. Assuming plane-stress in
XY"""
eps = 0.5 * (ufl.grad(u) + ufl.grad(u).T)
sigma = E / (1.0 - nu**2) * ((1.0 - nu) * eps + nu * ufl.Identity(2) * ufl.tr(eps))
return sigma
With the definitions above, we can define the different blocks of the variational formulation
To generate (C-code) and JIT compile the kernels, we use
ffcx_jit for each individual block.
We extract the kernel function from the compiled form object by
getting the tabulate_tensor_{dtype} attribute of the compiled form.
ufcx00, _, _ = ffcx_jit(msh.comm, a00, form_compiler_options={"scalar_type": dtype}) # type: ignore
kernel00 = getattr(ufcx00.form_integrals[0], f"tabulate_tensor_{np.dtype(dtype).name}") # type: ignore
ufcx01, _, _ = ffcx_jit(msh.comm, a01, form_compiler_options={"scalar_type": dtype}) # type: ignore
kernel01 = getattr(ufcx01.form_integrals[0], f"tabulate_tensor_{np.dtype(dtype).name}") # type: ignore
ufcx10, _, _ = ffcx_jit(msh.comm, a10, form_compiler_options={"scalar_type": dtype}) # type: ignore
kernel10 = getattr(ufcx10.form_integrals[0], f"tabulate_tensor_{np.dtype(dtype).name}") # type: ignore
ffi = cffi.FFI()
if np.issubdtype(dtype, np.complexfloating):
if cffi.__version_info__ > (1, 16, 99) and cffi.__version_info__ <= (1, 17, 1):
print(
"CFFI 1.17.0 and 1.17.1 has a bug for complex type."
"See https://github.com/FEniCS/dolfinx/pull/3635. Exiting."
)
exit(0)
cffi_support.register_type(ffi.typeof("double _Complex"), numba.types.complex128)
Get local dofmap sizes for later local tensor tabulations
Ssize = S.element.space_dimension
Usize = U.element.space_dimension
Next, we define a static condensation kernel that uses the
previously defined kernels to compute the condensed local element
tensor. The kernel is decorated with numba.cfunc() using the
appropriate signature obtained from ufcx_signature().`
@numba.cfunc(ufcx_signature(dtype, rtype), nopython=True) # type: ignore
def tabulate_A(A_, w_, c_, coords_, entity_local_index, permutation=ffi.NULL, custom_data=None):
"""Element kernel that applies static condensation."""
# Prepare target condensed local element tensor
A = numba.carray(A_, (Usize, Usize), dtype=dtype)
# Tabulate all sub blocks locally
A00 = np.zeros((Ssize, Ssize), dtype=dtype)
kernel00(
ffi.from_buffer(A00),
w_,
c_,
coords_,
entity_local_index,
permutation,
empty_void_pointer(),
)
A01 = np.zeros((Ssize, Usize), dtype=dtype)
kernel01(
ffi.from_buffer(A01),
w_,
c_,
coords_,
entity_local_index,
permutation,
empty_void_pointer(),
)
A10 = np.zeros((Usize, Ssize), dtype=dtype)
kernel10(
ffi.from_buffer(A10),
w_,
c_,
coords_,
entity_local_index,
permutation,
empty_void_pointer(),
)
# A = - A10 * A00^{-1} * A01
A[:, :] = -A10 @ np.linalg.solve(A00, A01)
Prepare a Form with a condensed
tabulation kernel. We specify the integration domains to be the
cells owned by the current process
formtype = form_cpp_class(dtype) # type: ignore
cells = np.arange(msh.topology.index_map(msh.topology.dim).size_local)
integrals = {IntegralType.cell: [(0, tabulate_A.address, cells, np.array([], dtype=np.int8))]}
a_cond = Form(
formtype([U._cpp_object, U._cpp_object], integrals, [], [], False, [], mesh=msh._cpp_object)
)
Next, we pass the compiled kernel to the standard assemble_matrix function to assemble
to the global condensed stiffness matrix. We also assemble the right-hand
side vector using assemble_vector and apply the boundary conditions by
applying lifting and
set bc.
A_cond = assemble_matrix(a_cond, bcs=[bc])
A_cond.assemble()
b = assemble_vector(b1)
apply_lifting(b, [a_cond], bcs=[[bc]])
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) # type: ignore
bc.set(b)
We use a PETSc.KSP solver to solve the
condensed linear system. The solution is stored in a
Function, while we pass the
underlying data wrapped as a PETSc.Vec
to the solver by calling petsc_vec on the
vector attribute of the function.
We verify the condensed solution by comparing against a standard, pure displacement based formulation
a = form(-ufl.inner(sigma_u(u), ufl.grad(v)) * ufl.dx)
A = assemble_matrix(a, bcs=[bc])
A.assemble()
Create BoundingBoxTree
using bb_tree constructor
for efficient computation of the ownership of a set of evaluation points
bb_tree = geometry.bb_tree(msh, tdim, padding=0.0)
Check against standard table value
p = np.array([[48.0, 52.0, 0.0]], dtype=np.float64)
cell_candidates = geometry.compute_collisions_points(bb_tree, p)
cells = geometry.compute_colliding_cells(msh, cell_candidates, p).array
uc.x.scatter_forward()
if len(cells) > 0:
value = uc.eval(p, cells[0])
print(value[1])
assert np.isclose(value[1], 23.95, rtol=1.0e-2)
Check the equality of displacement based and mixed condensed global matrices, i.e. check that condensation is exact
assert np.isclose((A - A_cond).norm(), 0.0)