# Static condensation of linear elasticity

Copyright (C) 2020 Michal Habera and Andreas Zilian

Download sources

- `Python script`

- `Jupyter notebook`

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

    u_bc = Function(U)
    u_bc.x.array[:] = 0

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

    a00 = ufl.inner(sigma, tau) * ufl.dx
    a10 = -ufl.inner(sigma, ufl.grad(v)) * ufl.dx
    a01 = -ufl.inner(sigma_u(u), tau) * ufl.dx

    f = ufl.as_vector([0.0, 1.0 / 16])
    b1 = form(-ufl.inner(f, v) * ds(1), dtype=dtype)  # type: ignore

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.

    uc = Function(U, name="u_from_condensation")
    solver = PETSc.KSP().create(A_cond.getComm())  # type: ignore
    solver.setOperators(A_cond)
    solver.solve(b, uc.x.petsc_vec)
    solver.destroy()

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)