MixedInterpolator

Author: pbrubeck

firedrake/assemble.py

@@ -18,7 +18,6 @@
 import finat.ufl
 from firedrake import (extrusion_utils as eutils, matrix, parameters, solving,
                        tsfc_interface, utils)
-from firedrake.formmanipulation import split_form
 from firedrake.adjoint_utils import annotate_assemble
 from firedrake.ufl_expr import extract_unique_domain
 from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit
@@ -570,36 +569,9 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             rank = len(expr.arguments())
             if rank > 2:
                 raise ValueError("Cannot assemble an Interpolate with more than two arguments")
-            # If argument numbers have been swapped => Adjoint.
-            arg_operand = ufl.algorithms.extract_arguments(operand)
-            is_adjoint = (arg_operand and arg_operand[0].number() == 0)
-
             # Get the target space
             V = v.function_space().dual()
 
-            # Dual interpolation from mixed source
-            if is_adjoint and len(V) > 1:
-                cur = 0
-                sub_operands = []
-                components = numpy.reshape(operand, (-1,))
-                for Vi in V:
-                    sub_operands.append(ufl.as_tensor(components[cur:cur+Vi.value_size].reshape(Vi.value_shape)))
-                    cur += Vi.value_size
-
-                # Component-split of the primal operands interpolated into the dual argument-split
-                split_interp = sum(reconstruct_interp(sub_operands[i], v=vi) for (i,), vi in split_form(v))
-                return assemble(split_interp, tensor=tensor)
-
-            # Dual interpolation into mixed target
-            if is_adjoint and len(arg_operand[0].function_space()) > 1 and rank == 1:
-                V = arg_operand[0].function_space()
-                tensor = tensor or firedrake.Cofunction(V.dual())
-
-                # Argument-split of the Interpolate gets assembled into the corresponding sub-tensor
-                for (i,), sub_interp in split_form(expr):
-                    assemble(sub_interp, tensor=tensor.subfunctions[i])
-                return tensor
-
             # Get the interpolator
             interp_data = expr.interp_data.copy()
             default_missing_val = interp_data.pop('default_missing_val', None)

firedrake/interpolation.py

@@ -263,7 +263,18 @@ class Interpolator(abc.ABC):
 
     def __new__(cls, expr, V, **kwargs):
         if isinstance(expr, ufl.Interpolate):
+            # Mixed spaces are handled well only by the primal 1-form.
+            # Are we a 2-form or a dual 1-form?
+            arguments = expr.arguments()
+            if any(not isinstance(a, Coargument) for a in arguments):
+                # Do we have mixed source or target spaces?
+                spaces = [a.function_space() for a in arguments]
+                if len(spaces) < 2:
+                    spaces.append(V)
+                if any(len(space) > 1 for space in spaces):
+                    return object.__new__(MixedInterpolator)
             expr, = expr.ufl_operands
+
         target_mesh = as_domain(V)
         source_mesh = extract_unique_domain(expr) or target_mesh
         submesh_interp_implemented = \
@@ -369,7 +380,7 @@ def _interpolate(self, *args, **kwargs):
         """
         pass
 
-    def assemble(self, tensor=None, default_missing_val=None):
+    def assemble(self, tensor=None, **kwargs):
         """Assemble the operator (or its action)."""
         from firedrake.assemble import assemble
         needs_adjoint = self.ufl_interpolate_renumbered != self.ufl_interpolate
@@ -383,13 +394,11 @@ def assemble(self, tensor=None, default_missing_val=None):
             if needs_adjoint:
                 # Out-of-place Hermitian transpose
                 petsc_mat.hermitianTranspose(out=res)
-            elif res:
-                petsc_mat.copy(res)
+            elif tensor:
+                petsc_mat.copy(tensor.petscmat)
             else:
                 res = petsc_mat
-            if tensor is None:
-                tensor = firedrake.AssembledMatrix(arguments, self.bcs, res)
-            return tensor
+            return tensor or firedrake.AssembledMatrix(arguments, self.bcs, res)
         else:
             # Assembling the action
             cofunctions = ()
@@ -401,11 +410,11 @@ def assemble(self, tensor=None, default_missing_val=None):
                     cofunctions = (dual_arg,)
 
             if needs_adjoint and len(arguments) == 0:
-                Iu = self._interpolate(default_missing_val=default_missing_val)
+                Iu = self._interpolate(**kwargs)
                 return assemble(ufl.Action(*cofunctions, Iu), tensor=tensor)
             else:
                 return self._interpolate(*cofunctions, output=tensor, adjoint=needs_adjoint,
-                                         default_missing_val=default_missing_val)
+                                         **kwargs)
 
 
 class DofNotDefinedError(Exception):
@@ -975,33 +984,10 @@ def callable():
         return callable
     else:
         loops = []
-        if len(V) == 1:
-            expressions = (expr,)
-        else:
-            if (hasattr(operand, "subfunctions") and len(operand.subfunctions) == len(V)
-                    and all(sub_op.ufl_shape == Vsub.value_shape for Vsub, sub_op in zip(V, operand.subfunctions))):
-                # Use subfunctions if they match the target shapes
-                operands = operand.subfunctions
-            else:
-                # Unflatten the expression into the shapes of the mixed components
-                offset = 0
-                operands = []
-                for Vsub in V:
-                    if len(Vsub.value_shape) == 0:
-                        operands.append(operand[offset])
-                    else:
-                        components = [operand[offset + j] for j in range(Vsub.value_size)]
-                        operands.append(ufl.as_tensor(numpy.reshape(components, Vsub.value_shape)))
-                    offset += Vsub.value_size
-
-            # Split the dual argument
-            if isinstance(dual_arg, Cofunction):
-                duals = dual_arg.subfunctions
-            elif isinstance(dual_arg, Coargument):
-                duals = [Coargument(Vsub, number=dual_arg.number()) for Vsub in dual_arg.function_space()]
-            else:
-                duals = [v for _, v in sorted(firedrake.formmanipulation.split_form(dual_arg))]
-            expressions = map(expr._ufl_expr_reconstruct_, operands, duals)
+        expressions = split_interpolate_target(expr)
+
+        if access == op2.INC:
+            loops.append(tensor.zero)
 
         # Interpolate each sub expression into each function space
         for Vsub, sub_tensor, sub_expr in zip(V, tensor, expressions):
@@ -1074,8 +1060,6 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     parameters['scalar_type'] = utils.ScalarType
 
     callables = ()
-    if access == op2.INC:
-        callables += (tensor.zero,)
 
     # For the matfree adjoint 1-form and the 0-form, the cellwise kernel will add multiple
     # contributions from the facet DOFs of the dual argument.
@@ -1720,3 +1704,115 @@ def _wrap_dummy_mat(self):
 
     def duplicate(self, mat=None, op=None):
         return self._wrap_dummy_mat()
+
+
+def split_interpolate_target(expr: ufl.Interpolate):
+    """Split an Interpolate into the components (subfunctions) of the target space."""
+    dual_arg, operand = expr.argument_slots()
+    V = dual_arg.function_space().dual()
+    if len(V) == 1:
+        return (expr,)
+    # Split the target (dual) argument
+    if isinstance(dual_arg, Cofunction):
+        duals = dual_arg.subfunctions
+    elif isinstance(dual_arg, ufl.Coargument):
+        duals = [Coargument(Vsub, dual_arg.number()) for Vsub in dual_arg.function_space()]
+    else:
+        duals = [vi for _, vi in sorted(firedrake.formmanipulation.split_form(dual_arg))]
+    # Split the operand into the target shapes
+    if (isinstance(operand, firedrake.Function) and len(operand.subfunctions) == len(V)
+            and all(fsub.ufl_shape == Vsub.value_shape for Vsub, fsub in zip(V, operand.subfunctions))):
+        # Use subfunctions if they match the target shapes
+        operands = operand.subfunctions
+    else:
+        # Unflatten the expression into the target shapes
+        cur = 0
+        operands = []
+        components = numpy.reshape(operand, (-1,))
+        for Vi in V:
+            operands.append(ufl.as_tensor(components[cur:cur+Vi.value_size].reshape(Vi.value_shape)))
+            cur += Vi.value_size
+    expressions = tuple(map(expr._ufl_expr_reconstruct_, operands, duals))
+    return expressions
+
+
+class MixedInterpolator(Interpolator):
+    """A reusable interpolation object between MixedFunctionSpaces.
+
+    Parameters
+    ----------
+    expr
+        The underlying ufl.Interpolate or the operand to the ufl.Interpolate.
+    V
+        The :class:`.FunctionSpace` or :class:`.Function` to
+        interpolate into.
+    bcs
+        A list of boundary conditions.
+    **kwargs
+        Any extra kwargs are passed on to the sub Interpolators.
+        For details see :class:`firedrake.interpolation.Interpolator`.
+    """
+    def __init__(self, expr, V, bcs=None, **kwargs):
+        if not isinstance(expr, ufl.Interpolate):
+            fs = V if isinstance(V, ufl.FunctionSpace) else V.function_space()
+            expr = interpolate(expr, fs)
+        if bcs is None:
+            bcs = ()
+        self.expr = expr
+        self.V = V
+        self.bcs = bcs
+        self.arguments = expr.arguments()
+
+        # Split the target (dual) argument
+        dual_split = split_interpolate_target(expr)
+        self.sub_interpolators = {}
+        for i, form in enumerate(dual_split):
+            # Split the source (primal) argument
+            for j, sub_interp in firedrake.formmanipulation.split_form(form):
+                j = max(j) if j else 0
+                # Ensure block sparsity
+                vi, operand = sub_interp.argument_slots()
+                if not isinstance(operand, ufl.classes.Zero):
+                    Vtarget = vi.function_space().dual()
+                    adjoint = vi.number() == 1 if isinstance(vi, Coargument) else True
+
+                    args = sub_interp.arguments()
+                    Vsource = args[0 if adjoint else 1].function_space()
+                    sub_bcs = [bc for bc in bcs if bc.function_space() in {Vsource, Vtarget}]
+
+                    indices = (j, i) if adjoint else (i, j)
+                    Isub = Interpolator(sub_interp, Vtarget, bcs=sub_bcs, **kwargs)
+                    self.sub_interpolators[indices] = Isub
+
+    def assemble(self, tensor=None, **kwargs):
+        """Assemble the operator (or its action)."""
+        rank = len(self.arguments)
+        if rank == 2:
+            # Assemble the operator
+            sub_tensors = {}
+            for ij, Isub in self.sub_interpolators.items():
+                block = tensor.petscmat.getNestSubMatrix(*ij) if tensor else PETSc.Mat()
+                sub_tensors[ij] = firedrake.AssembledMatrix(Isub.arguments, Isub.bcs, block)
+                Isub.assemble(tensor=sub_tensors[ij], **kwargs)
+            if tensor is None:
+                shape = tuple(len(a.function_space()) for a in self.arguments)
+                blocks = numpy.reshape([sub_tensors[ij].petscmat if ij in sub_tensors else PETSc.Mat()
+                                        for ij in numpy.ndindex(shape)], shape)
+                petscmat = PETSc.Mat().createNest(blocks)
+                tensor = firedrake.AssembledMatrix(self.arguments, self.bcs, petscmat)
+        elif rank == 1:
+            # Assemble the action
+            if tensor is None:
+                V_dest = self.arguments[0].function_space().dual()
+                tensor = firedrake.Function(V_dest)
+            for k, fsub in enumerate(tensor.subfunctions):
+                fsub.assign(sum(Isub.assemble(**kwargs) for (i, j), Isub in self.sub_interpolators.items() if i == k))
+        elif rank == 0:
+            # Assemble the double action
+            result = sum(Isub.assemble(**kwargs) for (i, j), Isub in self.sub_interpolators.items())
+            return tensor.assign(result) if tensor else result
+        return tensor
+
+    def _interpolate(self, output=None, **kwargs):
+        """Assemble the action."""
+        return self.assemble(tensor=output, **kwargs)

tests/firedrake/regression/test_interpolate.py

@@ -517,3 +517,50 @@ def test_interpolate_logical_not():
     a = assemble(interpolate(conditional(Not(x < .2), 1, 0), V))
     b = assemble(interpolate(conditional(x >= .2, 1, 0), V))
     assert np.allclose(a.dat.data, b.dat.data)
+
+
+@pytest.mark.parametrize("mode", ("forward", "adjoint"))
+def test_mixed_matrix(mode):
+    nx = 3
+    mesh = UnitSquareMesh(nx, nx)
+
+    V1 = VectorFunctionSpace(mesh, "CG", 2)
+    V2 = FunctionSpace(mesh, "CG", 1)
+    V3 = FunctionSpace(mesh, "CG", 1)
+    V4 = FunctionSpace(mesh, "DG", 1)
+
+    Z = V1 * V2
+    W = V3 * V3 * V4
+
+    if mode == "forward":
+        I = Interpolate(TrialFunction(Z), TestFunction(W.dual()))
+        a = assemble(I)
+        assert a.arguments()[0].function_space() == W.dual()
+        assert a.arguments()[1].function_space() == Z
+        assert a.petscmat.getSize() == (W.dim(), Z.dim())
+        assert a.petscmat.getType() == "nest"
+
+        u = Function(Z)
+        u.subfunctions[0].sub(0).assign(1)
+        u.subfunctions[0].sub(1).assign(2)
+        u.subfunctions[1].assign(3)
+        result_matfree = assemble(Interpolate(u, TestFunction(W.dual())))
+    elif mode == "adjoint":
+        I = Interpolate(TestFunction(Z), TrialFunction(W.dual()))
+        a = assemble(I)
+        assert a.arguments()[1].function_space() == W.dual()
+        assert a.arguments()[0].function_space() == Z
+        assert a.petscmat.getSize() == (Z.dim(), W.dim())
+        assert a.petscmat.getType() == "nest"
+
+        u = Function(W.dual())
+        u.subfunctions[0].assign(1)
+        u.subfunctions[1].assign(2)
+        u.subfunctions[2].assign(3)
+        result_matfree = assemble(Interpolate(TestFunction(Z), u))
+    else:
+        raise ValueError(f"Unrecognized mode {mode}")
+
+    result_explicit = assemble(action(a, u))
+    for x, y in zip(result_explicit.subfunctions, result_matfree.subfunctions):
+        assert np.allclose(x.dat.data, y.dat.data)