Assemble mixed cross-mesh interpolation operator

docs/source/interpolation.rst

@@ -109,7 +109,7 @@ This also works when interpolating into a space defined on the facets of the mes
 
 
 Semantics of symbolic interpolation
---------------------------------------------
+-----------------------------------
 
 Let :math:`U` and :math:`V` be finite element spaces with DoFs :math:`\{\psi^{*}_{i}\}` and :math:`\{\phi^{*}_{i}\}`
 and basis functions :math:`\{\psi_{i}\}` and :math:`\{\phi_{i}\}`, respectively.
@@ -430,6 +430,63 @@ the next operation using ``f``.
 For interaction with external point data, see the
 :ref:`corresponding manual section <external-point-data>`.
 
+Interpolation between mixed function spaces
+-------------------------------------------
+
+Assembly of interpolation operators between mixed function spaces is also supported.
+Each component of the mixed space may be on different meshes.
+For example, consider the following mixed finite element spaces:
+
+.. math::
+
+   W &= V_1 \times V_2 \\
+   U &= V_3 \times V_4
+
+where each :math:`V_i` is a finite element space defined on possibly different meshes.
+We can assemble the interpolation matrix from :math:`U` to :math:`W` in Firedrake as follows:
+
+.. literalinclude:: ../../tests/firedrake/regression/test_interpolation_manual.py
+   :language: python3
+   :dedent:
+   :start-after: [test_mixed_space_interpolation 1]
+   :end-before: [test_mixed_space_interpolation 2]
+
+We specified ``mat_type="nest"`` here to obtain a PETSc MatNest matrix, but Firedrake also
+supports assembly of ``mat_type="aij"`` and ``mat_type="matfree"`` interpolation matrices
+between mixed function spaces. In this example ``I`` is a block diagonal matrix, with
+each block given by
+
+.. math::
+
+   \begin{pmatrix}
+   V_3 \rightarrow V_1 & 0 \\
+   0 & V_4 \rightarrow V_2
+   \end{pmatrix}
+
+The off-diagonal blocks are zero since the dofs are applied component-wise. Firedrake's form
+compiler recognises this and avoids assembling the zero blocks.
+
+We can assemble more general interpolation matrices between mixed function spaces by interpolating
+vector expressions with arguments. For example, by doing
+
+.. literalinclude:: ../../tests/firedrake/regression/test_interpolation_manual.py
+   :language: python3
+   :dedent:
+   :start-after: [test_mixed_space_interpolation 3]
+   :end-before: [test_mixed_space_interpolation 4]
+
+we can assemble the interpolation matrix with block structure
+
+.. math::
+
+   \begin{pmatrix}
+   V_3 \rightarrow V_1 & V_4 \rightarrow V_1 \\
+   V_3 \rightarrow V_2 & V_4 \rightarrow V_2
+   \end{pmatrix}
+
+Here we obtain non-zero off-diagonal blocks by including both components of the trial function
+in each component of the expression.
+
 Generating Functions with randomised values
 -------------------------------------------
 

firedrake/interpolation.py

@@ -257,7 +257,12 @@ def __init__(self, expr: Interpolate):
         # Delay calling .unique() because MixedInterpolator is fine with MeshSequence
         self.target_mesh = self.target_space.mesh()
         """The domain we are interpolating into."""
-        self.source_mesh = extract_unique_domain(operand) or self.target_mesh
+
+        try:
+            source_mesh = extract_unique_domain(operand)
+        except ValueError:
+            source_mesh = extract_unique_domain(operand, expand_mesh_sequence=False)
+        self.source_mesh = source_mesh or self.target_mesh
         """The domain we are interpolating from."""
 
         # Interpolation options
@@ -416,7 +421,6 @@ class CrossMeshInterpolator(Interpolator):
     @no_annotations
     def __init__(self, expr: Interpolate):
         super().__init__(expr)
-        self.target_mesh = self.target_mesh.unique()
         if self.access and self.access != op2.WRITE:
             raise NotImplementedError(
                 "Access other than op2.WRITE not implemented for cross-mesh interpolation."
@@ -440,7 +444,7 @@ def __init__(self, expr: Interpolate):
         else:
             self.missing_points_behaviour = MissingPointsBehaviour.ERROR
 
-        if self.source_mesh.geometric_dimension != self.target_mesh.geometric_dimension:
+        if self.source_mesh.unique().geometric_dimension != self.target_mesh.unique().geometric_dimension:
             raise ValueError("Geometric dimensions of source and destination meshes must match.")
 
         dest_element = self.target_space.ufl_element()
@@ -477,12 +481,12 @@ def _get_symbolic_expressions(self) -> tuple[Interpolate, Interpolate]:
         """
         from firedrake.assemble import assemble
         # Immerse coordinates of target space point evaluation dofs in src_mesh
-        target_space_vec = VectorFunctionSpace(self.target_mesh, self.dest_element)
-        f_dest_node_coords = assemble(interpolate(self.target_mesh.coordinates, target_space_vec))
-        dest_node_coords = f_dest_node_coords.dat.data_ro.reshape(-1, self.target_mesh.geometric_dimension)
+        target_space_vec = VectorFunctionSpace(self.target_mesh.unique(), self.dest_element)
+        f_dest_node_coords = assemble(interpolate(self.target_mesh.unique().coordinates, target_space_vec))
+        dest_node_coords = f_dest_node_coords.dat.data_ro.reshape(-1, self.target_mesh.unique().geometric_dimension)
         try:
             vom = VertexOnlyMesh(
-                self.source_mesh,
+                self.source_mesh.unique(),
                 dest_node_coords,
                 redundant=False,
                 missing_points_behaviour=self.missing_points_behaviour,
@@ -607,11 +611,10 @@ class SameMeshInterpolator(Interpolator):
     @no_annotations
     def __init__(self, expr):
         super().__init__(expr)
-        self.target_mesh = self.target_mesh.unique()
         subset = self.subset
         if subset is None:
-            target = self.target_mesh.topology
-            source = self.source_mesh.topology
+            target = self.target_mesh.unique().topology
+            source = self.source_mesh.unique().topology
             if all(isinstance(m, MeshTopology) for m in [target, source]) and target is not source:
                 composed_map, result_integral_type = source.trans_mesh_entity_map(target, "cell", "everywhere", None)
                 if result_integral_type != "cell":
@@ -652,7 +655,7 @@ def _get_tensor(self, mat_type: Literal["aij", "baij"]) -> op2.Mat | Function |
             The tensor to interpolate into.
         """
         if self.rank == 0:
-            R = FunctionSpace(self.target_mesh, "Real", 0)
+            R = FunctionSpace(self.target_mesh.unique(), "Real", 0)
             f = Function(R, dtype=ScalarType)
         elif self.rank == 1:
             f = Function(self.ufl_interpolate.function_space())
@@ -689,8 +692,8 @@ def _get_monolithic_sparsity(self, mat_type: Literal["aij", "baij"]) -> op2.Spar
         Vcol = self.interpolate_args[1].function_space()
         if len(Vrow) > 1 or len(Vcol) > 1:
             raise NotImplementedError("Interpolation matrix with MixedFunctionSpace requires MixedInterpolator")
-        Vrow_map = get_interp_node_map(self.source_mesh, self.target_mesh, Vrow)
-        Vcol_map = get_interp_node_map(self.source_mesh, self.target_mesh, Vcol)
+        Vrow_map = get_interp_node_map(self.source_mesh.unique(), self.target_mesh.unique(), Vrow)
+        Vcol_map = get_interp_node_map(self.source_mesh.unique(), self.target_mesh.unique(), Vcol)
         sparsity = op2.Sparsity((Vrow.dof_dset, Vcol.dof_dset),
                                 [(Vrow_map, Vcol_map, None)],  # non-mixed
                                 name=f"{Vrow.name}_{Vcol.name}_sparsity",
@@ -1689,7 +1692,7 @@ def _build_aij(
 
     def _get_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None):
         mat_type = mat_type or "aij"
-        sub_mat_type = sub_mat_type or "baij"
+        sub_mat_type = sub_mat_type or "aij"
         Isub = self._get_sub_interpolators(bcs=bcs)
         V_dest = self.ufl_interpolate.function_space() or self.target_space
         f = tensor or Function(V_dest)

tests/firedrake/regression/test_interpolate_cross_mesh.py

@@ -1,6 +1,8 @@
 from firedrake import *
 from firedrake.petsc import DEFAULT_PARTITIONER
 from firedrake.ufl_expr import extract_unique_domain
+from firedrake.mesh import Mesh, plex_from_cell_list
+from firedrake.formmanipulation import split_form
 import numpy as np
 import pytest
 from ufl import product
@@ -614,8 +616,8 @@ def test_line_integral():
     # Create a 1D line mesh in 2D from (0, 0) to (1, 1) with 1 cell
     cells = np.asarray([[0, 1]])
     vertex_coords = np.asarray([[0.0, 0.0], [1.0, 1.0]])
-    plex = mesh.plex_from_cell_list(1, cells, vertex_coords, comm=m.comm)
-    line = mesh.Mesh(plex, dim=2)
+    plex = plex_from_cell_list(1, cells, vertex_coords, comm=m.comm)
+    line = Mesh(plex, dim=2)
     x, y = SpatialCoordinate(line)
     V_line = FunctionSpace(line, "CG", 2)
     f_line = Function(V_line).interpolate(x * y)
@@ -624,8 +626,8 @@ def test_line_integral():
     # Create a 1D line around the unit square (2D) with 4 cells
     cells = np.asarray([[0, 1], [1, 2], [2, 3], [3, 0]])
     vertex_coords = np.asarray([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])
-    plex = mesh.plex_from_cell_list(1, cells, vertex_coords, comm=m.comm)
-    line_square = mesh.Mesh(plex, dim=2)
+    plex = plex_from_cell_list(1, cells, vertex_coords, comm=m.comm)
+    line_square = Mesh(plex, dim=2)
     x, y = SpatialCoordinate(line_square)
     V_line_square = FunctionSpace(line_square, "CG", 2)
     f_line_square = Function(V_line_square).interpolate(x * y)
@@ -751,3 +753,40 @@ def test_interpolate_cross_mesh_interval(periodic):
     f_dest = Function(V_dest).interpolate(f_src)
     x_dest, = SpatialCoordinate(m_dest)
     assert abs(assemble((f_dest - (-(x_dest - .5) ** 2)) ** 2 * dx)) < 1.e-16
+
+
+def test_mixed_interpolator_cross_mesh():
+    # Tests assembly of mixed interpolator across meshes
+    mesh1 = UnitSquareMesh(4, 4)
+    mesh2 = UnitSquareMesh(3, 3, quadrilateral=True)
+    mesh3 = UnitDiskMesh(2)
+    mesh4 = UnitTriangleMesh(3)
+    V1 = FunctionSpace(mesh1, "CG", 1)
+    V2 = FunctionSpace(mesh2, "CG", 2)
+    V3 = FunctionSpace(mesh3, "CG", 3)
+    V4 = FunctionSpace(mesh4, "CG", 4)
+
+    W = V1 * V2
+    U = V3 * V4
+
+    w = TrialFunction(W)
+    w0, w1 = split(w)
+    expr = as_vector([w0 + w1, w0 + w1])
+    mixed_interp = interpolate(expr, U, allow_missing_dofs=True)  # Interpolating from W to U
+
+    # The block matrix structure is
+    # | V1 -> V3   V2 -> V3 |
+    # | V1 -> V4   V2 -> V4 |
+
+    res = assemble(mixed_interp, mat_type="nest")
+    assert isinstance(res, AssembledMatrix)
+    assert res.petscmat.type == "nest"
+
+    split_interp = dict(split_form(mixed_interp))
+
+    for i in range(2):
+        for j in range(2):
+            interp_ij = split_interp[(i, j)]
+            assert isinstance(interp_ij, Interpolate)
+            res_block = assemble(interpolate(TrialFunction(W.sub(j)), U.sub(i), allow_missing_dofs=True))
+            assert np.allclose(res.petscmat.getNestSubMatrix(i, j)[:, :], res_block.petscmat[:, :])

tests/firedrake/regression/test_interpolation_manual.py

@@ -1,4 +1,5 @@
 from firedrake import *
+from firedrake.formmanipulation import split_form
 import pytest
 import numpy as np
 
@@ -235,3 +236,50 @@ def correct_indent():
 
     assert np.isclose(dest_eval05.evaluate(f_dest), 0.5)  # x_src^2 + y_src^2 = 0.5
     assert np.isclose(dest_eval15.evaluate(f_dest), 3.0)  # x_dest + y_dest = 3.0
+
+
+def test_mixed_space_interpolation():
+    mesh = UnitSquareMesh(2, 2)
+    V1 = FunctionSpace(mesh, "CG", 1)
+    V2 = FunctionSpace(mesh, "CG", 2)
+    V3 = FunctionSpace(mesh, "CG", 3)
+    V4 = FunctionSpace(mesh, "CG", 4)
+    W = V1 * V2
+    U = V3 * V4
+
+    # [test_mixed_space_interpolation 1]
+    interp = interpolate(TrialFunction(U), W)
+    I = assemble(interp, mat_type="nest")
+    # [test_mixed_space_interpolation 2]
+
+    # The block matrix structure is
+    # | V3 -> V1      0    |
+    # |     0     V4 -> V2 |
+    for i in range(2):
+        for j in range(2):
+            sub_mat = I.petscmat.getNestSubMatrix(i, j)
+            if i != j:
+                assert not sub_mat
+                continue
+            else:
+                res_block = assemble(interpolate(TrialFunction(U.sub(j)), W.sub(i)))
+                assert np.allclose(sub_mat[:, :], res_block.petscmat[:, :])
+                assert sub_mat.type == "seqaij"
+
+    # [test_mixed_space_interpolation 3]
+    u0, u1 = TrialFunctions(U)
+    expr = as_vector([u0 + u1, u0 + u1])
+    interp = interpolate(expr, W)
+    I2 = assemble(interp, mat_type="nest")
+    # [test_mixed_space_interpolation 4]
+
+    # The block matrix structure is
+    # | V3 -> V1   V4 -> V1 |
+    # | V3 -> V2   V4 -> V2 |
+    split_interp = dict(split_form(interp))
+    for i in range(2):
+        for j in range(2):
+            interp_ij = split_interp[(i, j)]
+            assert isinstance(interp_ij, Interpolate)
+            res_block = assemble(interpolate(TrialFunction(U.sub(j)), W.sub(i), allow_missing_dofs=True))
+            assert np.allclose(I2.petscmat.getNestSubMatrix(i, j)[:, :], res_block.petscmat[:, :])

tests/firedrake/regression/test_interpolator_types.py

@@ -163,8 +163,8 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type):
                 # Always seqaij for scalar
                 assert sub_mat.type == "seqaij"
             else:
-                # matnest sub_mat_type defaults to baij
-                assert sub_mat.type == "seq" + (sub_mat_type if sub_mat_type else "baij")
+                # matnest sub_mat_type defaults to aij
+                assert sub_mat.type == "seq" + (sub_mat_type if sub_mat_type else "aij")
 
     with pytest.raises(NotImplementedError):
         assemble(interp, mat_type="baij")