Add interpolation matrix of ufl.core.expr.Expr

README.md

@@ -18,6 +18,7 @@ The package is still in its early stages and many functionalities are still miss
 - Maps between degrees of freedom and vertices: `vertex_to_dofmap` and `dof_to_vertex`
 - Blocked Newton Solver
 - Function evaluation at specified points
+- Interpolation matrices from any `ufl.core.expr.Expr` into a compatible space.
 
 ## Installation
 

pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "scifem"
-version = "0.12.0"
+version = "0.13.0"
 description = "Scientific tools for finite element methods"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -138,7 +138,7 @@ tag = true
 sign_tags = false
 tag_name = "v{new_version}"
 tag_message = "Bump version: {current_version} → {new_version}"
-current_version = "0.12.0"
+current_version = "0.13.0"
 
 
 [[tool.bumpversion.files]]

src/scifem/__init__.py

@@ -20,6 +20,8 @@
     reverse_mark_entities,
 )
 from .eval import evaluate_function
+from .interpolation import interpolation_matrix, prepare_interpolation_data
+
 
 __all__ = [
     "PointSource",
@@ -43,9 +45,17 @@
     "reverse_mark_entities",
     "norm",
     "interpolate_function_onto_facet_dofs",
+    "interpolation_matrix",
+    "prepare_interpolation_data",
 ]
 
 
+if dolfinx.has_petsc4py:
+    from .interpolation import petsc_interpolation_matrix  # type: ignore
+
+    __all__ += ["petsc_interpolation_matrix"]
+
+
 def vertex_to_dofmap(V: dolfinx.fem.FunctionSpace) -> npt.NDArray[np.int32]:
     """
     Create a map from the vertices (local to the process) to the correspondning degrees

src/scifem/interpolation.py

@@ -0,0 +1,254 @@
+import dolfinx
+import ufl
+import numpy as np
+import numpy.typing as npt
+from .utils import unroll_dofmap
+
+__all__ = ["interpolation_matrix", "prepare_interpolation_data"]
+
+if dolfinx.has_petsc4py:
+    from petsc4py import PETSc
+
+    __all__.append("petsc_interpolation_matrix")
+
+
+def prepare_interpolation_data(
+    expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace
+) -> npt.NDArray[np.inexact]:
+    """Convenience function for preparing data required for assembling the interpolation matrix
+
+    .. math::
+        \\begin{align*}
+        \\Lambda: V &\\rightarrow Q \\\\
+        \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j
+        \\end{align*}
+
+    where :math:`l_j` is the dual basis of the space :math:`Q` with basis functions :math:`\\phi_j`,
+    and :math:`\\psi_j` are the basis functions of the space :math:`V`.
+
+    Args:
+        expr: The UFL expression containing a trial function from space `V`
+        Q: Output interpolation space
+    Returns:
+        Interpolation data per cell, as an numpy array.
+    """
+    try:
+        q_points = Q.element.interpolation_points()
+    except TypeError:
+        q_points = Q.element.interpolation_points
+
+    arguments = ufl.algorithms.extract_arguments(expr)
+    assert len(arguments) == 1
+    V = arguments[0].ufl_function_space()
+
+    num_points = q_points.shape[0]
+    compiled_expr = dolfinx.fem.Expression(expr, q_points)
+    mesh = Q.mesh
+    tdim = mesh.topology.dim
+    num_cells = mesh.topology.index_map(tdim).size_local
+    #
+    # (num_cells, num_points, num_dofs*bs, expr_value_size)
+    array_evaluated = compiled_expr.eval(mesh, np.arange(num_cells, dtype=np.int32))
+    assert np.prod(Q.value_shape) == np.prod(expr.ufl_shape)
+
+    im = Q.element.basix_element.interpolation_matrix
+
+    # Get data as (num_cells*num_points,1, expr_shape, num_test_basis_functions*test_block_size)
+    expr_size = int(np.prod(expr.ufl_shape))
+    array_evaluated = array_evaluated.reshape(
+        num_cells * q_points.shape[0], 1, expr_size, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
+    )
+    jacobian = dolfinx.fem.Expression(ufl.Jacobian(mesh), q_points)
+    detJ = dolfinx.fem.Expression(ufl.JacobianDeterminant(mesh), q_points)
+    K = dolfinx.fem.Expression(ufl.JacobianInverse(mesh), q_points)
+    jacs = jacobian.eval(mesh, np.arange(num_cells, dtype=np.int32)).reshape(
+        num_cells * num_points, mesh.geometry.dim, mesh.topology.dim
+    )
+    detJs = detJ.eval(mesh, np.arange(num_cells, dtype=np.int32)).flatten()
+    Ks = K.eval(mesh, np.arange(num_cells, dtype=np.int32)).reshape(
+        num_cells * num_points, mesh.geometry.dim, mesh.topology.dim
+    )
+
+    Q_vs = Q.element.basix_element.value_size
+    new_array = np.zeros(
+        (num_cells * num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs),
+        dtype=np.float64,
+    )
+    for i in range(V.dofmap.bs * V.dofmap.dof_layout.num_dofs):
+        for q in range(Q.dofmap.bs):
+            new_array[:, q * Q_vs : (q + 1) * Q_vs, i] = Q.element.basix_element.pull_back(
+                array_evaluated[:, :, q * Q_vs : (q + 1) * Q_vs, i], jacs, detJs, Ks
+            ).reshape(num_cells * num_points, Q_vs)
+    new_array = new_array.reshape(
+        num_cells, num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
+    )
+    interpolated_matrix = np.zeros(
+        (
+            num_cells,
+            Q.dofmap.dof_layout.num_dofs * Q.dofmap.bs,
+            V.dofmap.bs * V.dofmap.dof_layout.num_dofs,
+        ),
+        dtype=np.float64,
+    )
+
+    for c in range(num_cells):
+        for i in range(V.dofmap.bs * V.dofmap.dof_layout.num_dofs):
+            tmp_array = np.zeros((int(num_points), Q.dofmap.bs * Q_vs), dtype=np.float64)
+            for p in range(num_points):
+                tmp_array[p] = new_array[c, p, :, i]
+            if Q.dofmap.bs == 1:
+                interpolated_matrix[c, :, i] = (im @ tmp_array.T.flatten()).flatten()
+            else:
+                for q in range(Q.dofmap.bs):
+                    interpolated_matrix[c, q :: Q.dofmap.bs, i] = (
+                        im @ tmp_array.T[q].flatten()
+                    ).flatten()
+    # Apply dof transformation to each column (using Piopla maps)
+    mesh.topology.create_entity_permutations()
+    if Q.element.needs_dof_transformations:
+        cell_perm = mesh.topology.get_cell_permutation_info()[:num_cells]
+
+        permuted_matrix = interpolated_matrix.flatten().copy()
+        Q.element.Tt_inv_apply(
+            permuted_matrix, cell_perm, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
+        )
+    else:
+        permuted_matrix = interpolated_matrix.flatten()
+    return permuted_matrix.reshape(interpolated_matrix.shape)
+
+
+def interpolation_matrix(
+    expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace
+) -> dolfinx.la.MatrixCSR:
+    """Create the interpolation matrix :math:`\\Lambda` of a
+    :py:class:`UFL-expression<ufl.core.expr.Expr>` such that
+
+    .. math::
+        \\begin{align*}
+        \\Lambda: V &\\rightarrow Q \\\\
+        \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j
+        \\end{align*}
+
+    where :math:`l_j` is the dual basis of the space :math:`Q` with
+    basis functions :math:`\\phi_j`, and :math:`\\psi_j` are the basis functions of the
+    space :math:`V`.
+
+    Args:
+        expr: The UFL expression
+        Q: Output interpolation space
+
+    Returns:
+        Interpolation matrix as a :py:class:`MatrixCSR<dolfinx.la.MatrixCSR>`.
+    """
+
+    arguments = ufl.algorithms.extract_arguments(expr)
+    assert len(arguments) == 1
+    V = arguments[0].ufl_function_space()
+
+    interpolation_data = prepare_interpolation_data(expr, Q)
+
+    q = ufl.TestFunction(Q)
+    a = dolfinx.fem.form(ufl.inner(expr, q) * ufl.dx)
+
+    def scatter(
+        A: dolfinx.la.MatrixCSR,
+        num_cells: int,
+        dofs_visited: npt.NDArray[np.int32],
+        num_rows_local: int,
+        array_evaluated: npt.NDArray[np.inexact],
+        dofmap0: npt.NDArray[np.int32],
+        dofmap1: npt.NDArray[np.int32],
+    ):
+        A.data[:] = 0
+        for i in range(num_cells):
+            rows = dofmap0[i, :]
+            cols = dofmap1[i, :]
+            A_local = array_evaluated[i].reshape(len(rows), len(cols))
+            row_filter = (dofs_visited[rows] == 1) | (rows >= num_rows_local)
+            A_local[row_filter] = 0
+            A.add(A_local.flatten(), rows, cols)
+            dofs_visited[rows] = 1
+
+    A = dolfinx.fem.create_matrix(a)  # , dolfinx.la.BlockMode.expanded)
+
+    row_dofmap = unroll_dofmap(Q.dofmap.list, Q.dofmap.bs)  # (num_cells, num_rows)
+    col_dofmap = unroll_dofmap(V.dofmap.list, V.dofmap.bs)  # (num_cells, num_cols)
+
+    num_cells = Q.mesh.topology.index_map(Q.mesh.topology.dim).size_local
+    dofs_visited = np.zeros(
+        (Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts) * Q.dofmap.index_map_bs,
+        dtype=np.int8,
+    )
+    num_rows_local = Q.dofmap.index_map.size_local * Q.dofmap.bs
+    scatter(A, num_cells, dofs_visited, num_rows_local, interpolation_data, row_dofmap, col_dofmap)
+    A.scatter_reverse()
+    return A
+
+
+if dolfinx.has_petsc4py:
+
+    def petsc_interpolation_matrix(
+        expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace, use_petsc: bool = False
+    ) -> PETSc.Mat:
+        """Create the interpolation matrix :math:`\\Lambda` of a
+        :py:class:`UFL-expression<ufl.core.expr.Expr>` such that
+
+        .. math::
+            \\begin{align*}
+            \\Lambda: V &\\rightarrow Q \\\\
+            \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j
+            \\end{align*}
+
+        where :math:`l_j` is the dual basis of the space :math:`Q` with basis
+        functions :math:`\\phi_j`, and :math:`\\psi_j` are the basis functions
+        of the space :math:`V`.
+
+        Args:
+            expr: The UFL expression
+            Q: Output interpolation space
+
+        Returns:
+            Interpolation matrix as a :py:class:`PETSc.Mat<petsc4py.PETSc.Mat>`.
+        """
+        arguments = ufl.algorithms.extract_arguments(expr)
+        assert len(arguments) == 1
+        V = arguments[0].ufl_function_space()
+
+        interpolation_data = prepare_interpolation_data(expr, Q)
+
+        q = ufl.TestFunction(Q)
+        a = dolfinx.fem.form(ufl.inner(expr, q) * ufl.dx)
+        A = dolfinx.fem.petsc.create_matrix(a)
+
+        def scatter(
+            A: PETSc.Mat,
+            num_cells: int,
+            dofs_visited: npt.NDArray[np.int32],
+            num_rows_local: int,
+            array_evaluated: npt.NDArray[np.inexact],
+            dofmap0: npt.NDArray[np.int32],
+            dofmap1: npt.NDArray[np.int32],
+        ):
+            A.zeroEntries()
+            for i in range(num_cells):
+                rows = dofmap0[i, :]
+                cols = dofmap1[i, :]
+                A_local = array_evaluated[i].reshape(len(rows), len(cols))
+                row_filter = (dofs_visited[rows] == 1) | (rows >= num_rows_local)
+                A_local[row_filter] = 0
+                A.setValuesLocal(rows, cols, A_local, addv=PETSc.InsertMode.ADD_VALUES)
+                dofs_visited[rows] = 1
+
+        row_dofmap = unroll_dofmap(Q.dofmap.list, Q.dofmap.bs)  # (num_cells, num_rows)
+        col_dofmap = unroll_dofmap(V.dofmap.list, V.dofmap.bs)  # (num_cells, num_cols)
+        num_cells = Q.mesh.topology.index_map(Q.mesh.topology.dim).size_local
+        dofs_visited = np.zeros(
+            (Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts) * Q.dofmap.index_map_bs,
+            dtype=np.int8,
+        )
+        num_rows_local = Q.dofmap.index_map.size_local * Q.dofmap.bs
+        scatter(
+            A, num_cells, dofs_visited, num_rows_local, interpolation_data, row_dofmap, col_dofmap
+        )
+        A.assemble()
+        return A