Setting scalar function for BDM boundary condition with controlled flux [mixed Poisson]

To follow up on setting this strongly.
After some discussion with @cdaversin and @MiroK, I’ve reworked the strong imposition operator:

import numpy as np
import numpy.typing as npt

import ufl
import basix.ufl

from dolfinx import default_scalar_type as dolfinx_default_scalar_type
from dolfinx.fem import (
    FunctionSpace,
    Function,
    Expression,
)
from dolfinx.cpp.fem import CoordinateElement_float64
from dolfinx.mesh import compute_incident_entities


def strong_bc_bdm_function(
    Q: FunctionSpace,
    expr: ufl.core.expr.Expr,
    facets: npt.NDArray[np.int32],
) -> Function:
    """
    Create a function $u_h\in Q$ such that $u_h=\text{expr}$ for all dofs belonging
    to a subset of ``facets``. All other dofs are set to zero.

    Note:
        The resulting function  is only correct in the "normal" direction,
        i.e. :math:`u_{bc}\cdot n = expr`, while the tangential component is uncontrolled.
        This makes it hard to visualize the function when outputting it to file, either
        through interpolation to an appropriate DG space, or to a point-cloud.

    Args:
        Q: The function space to create the function $u_h$ in.
        expr: The expression to evaluate.
        facets: The facets on which to evaluate the expression.
    """
    domain = Q.mesh
    Q_el = Q.element
    fdim = domain.topology.dim - 1
    domain.topology.create_connectivity(fdim, domain.topology.dim)

    interpolation_points = Q_el.basix_element.x

    c_el = domain.ufl_domain().ufl_coordinate_element()
    ref_top = c_el.reference_topology
    ref_geom = c_el.reference_geometry
    facet_types = set(basix.cell.subentity_types(domain.basix_cell())[fdim])
    assert len(facet_types) == 1, "All facets must have the same topology"

    # Pull back interpolation points from reference coordinate element to facet reference element
    facet_cmap = basix.ufl.element(
        "Lagrange",
        facet_types.pop(),
        c_el.degree,
        shape=(domain.geometry.dim,),
        dtype=np.float64,
    )
    facet_cel = CoordinateElement_float64(facet_cmap.basix_element._e)
    reference_facet_points = None
    for i, points in enumerate(interpolation_points[fdim]):
        geom = ref_geom[ref_top[fdim][i]]

        ref_points = facet_cel.pull_back(points, geom)
        # Assert that interpolation points are all equal on all facets
        if reference_facet_points is None:
            reference_facet_points = ref_points
        else:
            assert np.allclose(reference_facet_points, ref_points)
    # Create expression for BC
    normal_expr = Expression(expr, reference_facet_points)

    points_per_entity = [sum(ip.shape[0] for ip in ips) for ips in interpolation_points]
    offsets = np.zeros(domain.topology.dim + 2, dtype=np.int32)
    offsets[1:] = np.cumsum(points_per_entity[: domain.topology.dim + 1])
    values_per_entity = np.zeros(
        (offsets[-1], domain.geometry.dim), dtype=dolfinx_default_scalar_type
    )

    # Compute integration entities (cell, local_facet index) for all facets
    all_connected_cells = compute_incident_entities(
        domain.topology, facets, domain.topology.dim - 1, domain.topology.dim
    )
    values = np.zeros(len(all_connected_cells) * offsets[-1] * domain.geometry.dim)
    domain.topology.create_connectivity(domain.topology.dim, fdim)
    c_to_f = domain.topology.connectivity(domain.topology.dim, fdim)
    num_facets_on_process = (
        domain.topology.index_map(fdim).size_local
        + domain.topology.index_map(fdim).num_ghosts
    )
    is_marked = np.zeros(num_facets_on_process, dtype=np.int8)
    is_marked[facets] = 1
    for i, cell in enumerate(all_connected_cells):
        values_per_entity[:] = 0.0
        local_facets = c_to_f.links(cell)
        for j, lf in enumerate(local_facets):
            if not is_marked[lf]:
                continue
            insert_pos = offsets[fdim] + reference_facet_points.shape[0] * j
            # Backwards compatibility
            entity = np.array([[cell, j]], dtype=np.int32)
            try:
                normal_on_facet = normal_expr.eval(domain, entity)
            except (AttributeError, AssertionError):
                normal_on_facet = normal_expr.eval(domain, entity.flatten())
            # NOTE: evaluate within loop to avoid large memory requirements
            values_per_entity[
                insert_pos : insert_pos + reference_facet_points.shape[0]
            ] = normal_on_facet.reshape(-1, domain.geometry.dim)
        values[
            i * offsets[-1] * domain.geometry.dim : (i + 1)
            * offsets[-1]
            * domain.geometry.dim
        ] = values_per_entity.reshape(-1)

    qh = Function(Q)
    qh._cpp_object.interpolate(
        values.reshape(-1, domain.geometry.dim).T.copy(), all_connected_cells
    )
    qh.x.scatter_forward()

    return qh

which now works with both v0.9.0 and the nightly branch, and added a note regarding visualizing the resulting BC function, as if you look at it, you need to remember that only the normal component is enforced, thus the tangential component at the boundaries will not be zero.
I also attach a code to illustrate this:

import ufl
import basix.ufl
import basix

from mpi4py import MPI
import numpy as np

from dolfinx.io import VTXWriter
from dolfinx.fem import (
    functionspace,
    Function,
)
from dolfinx.mesh import (
    create_unit_cube,
    exterior_facet_indices,
)

from solver import strong_bc_bdm_function

N = 4
mesh = create_unit_cube(MPI.COMM_WORLD, N, N, N)
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
tdim = mesh.topology.dim

order = 2
element_u = basix.ufl.element(basix.ElementFamily.BDM, mesh.basix_cell(), order)

V = functionspace(mesh, element_u)

regions_facets = exterior_facet_indices(mesh.topology)
nh = ufl.FacetNormal(mesh)
x = ufl.SpatialCoordinate(mesh)
inlet_expr = nh * (ufl.sin(x[0]) + x[1])
u_bc = strong_bc_bdm_function(V, inlet_expr, regions_facets)

from dolfinx.fem import locate_dofs_topological

bdofs = locate_dofs_topological(V, mesh.topology.dim - 1, regions_facets)


import basix.ufl
from dolfinx.fem import (
    Expression,
    IntegralType,
    functionspace,
    compute_integration_domains,
)
import dolfinx


def move_to_facet_quadrature(ufl_expr, mesh, sub_facets, scheme="default", degree=6):
    fdim = mesh.topology.dim - 1
    # Create submesh
    bndry_mesh, entity_map, _, _ = dolfinx.mesh.create_submesh(mesh, fdim, sub_facets)
    # Create quadrature space on submesh
    q_el = basix.ufl.quadrature_element(
        bndry_mesh.basix_cell(), ufl_expr.ufl_shape, scheme, degree
    )
    Q = functionspace(bndry_mesh, q_el)

    # Compute where to evaluate expression per submesh cell
    if hasattr(entity_map, "sub_topology_to_topology"):
        parent_entities = entity_map.sub_topology_to_topology(
            np.arange(bndry_mesh.topology.index_map(fdim).size_local, dtype=np.int32),
            inverse=False,
        )
        id_args = {}
    else:
        parent_entities = entity_map
        id_args = {"dim": fdim}
    integration_entities = compute_integration_domains(
        IntegralType.exterior_facet, mesh.topology, parent_entities, **id_args
    ).reshape(-1, 2)
    try:
        compiled_expr = Expression(ufl_expr, Q.element.interpolation_points)
    except AttributeError:
        compiled_expr = Expression(ufl_expr, Q.element.interpolation_points())

    # Evaluate expression
    q = Function(Q)
    try:
        q.x.array[:] = compiled_expr.eval(mesh, integration_entities).reshape(-1)
    except TypeError:
        q.x.array[:] = compiled_expr.eval(mesh, integration_entities.flatten()).reshape(
            -1
        )
    return q


exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
flux = u_bc
q = move_to_facet_quadrature(flux, mesh, exterior_facets, degree=8)
q.name = "u_at_quadrature"

import scifem

scifem.xdmf.create_pointcloud("flux.xdmf", [q])
V_DG = functionspace(mesh, ("DG", order + 1, (mesh.geometry.dim,)))
u_bc_DG = Function(V_DG)
u_bc_DG.interpolate(u_bc)
with VTXWriter(mesh.comm, "u_bc_vtx.bp", [u_bc_DG]) as vtx:
    vtx.write(0.0)


error = (
    ufl.inner(ufl.dot(u_bc, nh) * nh - inlet_expr, ufl.dot(u_bc, nh) * nh - inlet_expr)
    * ufl.ds
)

global_error = np.sqrt(
    mesh.comm.allreduce(
        dolfinx.fem.assemble_scalar(dolfinx.fem.form(error)), op=MPI.SUM
    )
)
print("L2 error in normal component of u_bc:", global_error, flush=True)

which yields

L2 error in normal component of u_bc: 6.667253143399306e-05

when the scaling functions is not in the space, and machine precision if it is.

I will probably add this function to scifem, with an illustrative demo in 2D soon.

3 Likes