Add facet mesh projector

[jorgensd]

class FacetMeshProjection:
    """Project an UFL expression from a mesh to a subset of its external facets
    
    :param mesh: The mesh the expression is defined on
    :param el: The element to project into on the facet space
    :param expr: The expression to project
    :param marker: The mesh tag marking the facets of the mesh
    :param marker_id: The id of the facets to use for the submesh
    :param quadrature_degree: Quadrature degree for integration
    :param petsc_options: PETSc options for linear solver
    """
    _mesh: dolfinx.mesh.Mesh # The mesh the velocity/pressure space is defined on
    _Q: dolfinx.fem.FunctionSpace # Quadrature space on boundary
    _problem: dolfinx.fem.petsc.LinearProblem # Linear problem for projection
    _solution: dolfinx.fem.Function # Solution function
    _facet_to_parent: np.ndarray # Mapping from facet to parent facet
    def __init__(
        self,
        mesh: dolfinx.mesh.Mesh,
        el: dolfinx.fem.ElementMetaData,
        expr: ufl.core.expr.Expr,
        marker: dolfinx.mesh.MeshTags,
        marker_id: typing.Tuple[int, ...],
        quadrature_degree: int,
        discontinuous:int = False,
        petsc_options: typing.Optional[dict] = None,
    ):
        if expr.ufl_domain() is not None:
            assert expr.ufl_domain() == mesh.ufl_domain()
        self._mesh = mesh
        
        if petsc_options is None:
            petsc_options = {}

        # Set up MeshView for boundary
        dim = marker.dim
        assert marker.dim == mesh.topology.dim-1
        all_facets = [marker.find(_id) for _id in marker_id]
        sub_facets = np.sort(np.hstack(all_facets))
        boundary_mesh, entity_map, _, _ = dolfinx.mesh.create_submesh(mesh, dim, sub_facets)
        self._facet_to_parent = entity_map        
        # Create quadrature space on boundary
        q_el = basix.ufl.quadrature_element(boundary_mesh.topology.cell_name(),
                                            value_shape=expr.ufl_shape , scheme="default", degree=quadrature_degree)


        self._Q = dolfinx.fem.functionspace(boundary_mesh, q_el)
        self._q = dolfinx.fem.Function(self._Q)
        # Build up array for parent (cell, local_facet_index pairs)
        mesh.topology.create_connectivity(dim, dim+1)
        mesh.topology.create_connectivity(dim+1, dim)
        f_to_c = mesh.topology.connectivity(dim, dim+1)
        c_to_f = mesh.topology.connectivity(dim+1, dim)
        self._parent_entities = np.empty(2*len(entity_map), dtype=np.int32)
        for i, entity in enumerate(entity_map):
            cells = f_to_c.links(entity)
            assert len(cells) == 1, "Facets in traction computation should be on a single cell."
            local_facets = c_to_f.links(cells[0])
            local_index = np.flatnonzero(local_facets == entity)[0]
            self._parent_entities[2*i] = cells[0]
            self._parent_entities[2*i+1] = local_index


        # Compile expression at facet quadrautre points
        self._compiled_expr = dolfinx.fem.Expression(expr, self._Q.element.interpolation_points())


        # Set up facet projection problem
        metadata = metadata={"quadrature_degree": quadrature_degree, "quadrature_scheme":"default"}
        facet_cell = dolfinx.cpp.mesh.cell_entity_type(mesh.topology.cell_type, mesh.topology.dim-1, 0)
        el = basix.ufl.element(el[0], facet_cell.name, el[1], shape=expr.ufl_shape, discontinuous=discontinuous,
                                dtype=dolfinx.default_real_type)
        self._V = dolfinx.fem.functionspace(boundary_mesh, el)
        self._solution = dolfinx.fem.Function(self._V)
        dx = ufl.dx(domain=boundary_mesh, metadata=metadata)
        u, v = ufl.TrialFunction(self._V), ufl.TestFunction(self._V)
        a =  ufl.inner(u, v) * dx
        L = ufl.inner(self._q, v) * dx
        self._problem = dolfinx.fem.petsc.LinearProblem(a, L, u=self._solution, petsc_options=petsc_options)
        self.assemble_lhs()

    def assemble_lhs(self):
        """Assemble lhs mass matrix"""
        self._problem.A.zeroEntries()
        dolfinx.fem.petsc.assemble_matrix_mat(self._problem.A, self._problem.a, bcs=self._problem.bcs)
        self._problem.A.assemble()

    def assemble_rhs(self):
        """Assemble rhs vector"""
        # Update expression at quadrature points

        values = self._compiled_expr.eval(self._mesh, self._parent_entities)
        self._q.x.array[:] = values.reshape(-1)
        # Assemble right hand side
        with self._problem.b.localForm() as b_loc:
            b_loc.set(0)
        dolfinx.fem.petsc.assemble_vector(self._problem.b, self._problem.L)
        self._problem.b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        self._problem.b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    
    def solve(self, assemble_lhs:bool=False, assemble_rhs:bool=True):
        """Project expression

        :param assemble_lhs: If `True` re-assemble the mass matrix, defaults to False
        :param assemble_rhs: If `True` re-assemble the right hand side, defaults to True
        """
        if assemble_lhs:
            self.assemble_lhs()
        if assemble_rhs:
            self.assemble_rhs()
        self._problem.solver.solve(self._problem.b, self._problem.u.vector)
        self._solution.x.scatter_forward()
        assert self._problem.solver.getConvergedReason() > 0, "Linear solver did not converge"

    @property
    def solution(self):
        return self._solution

    @property
    def parent_map(self):
        return self._facet_to_parent