Dual interpolation from/into MixedFunctionSpace (#4197)

* Interpolation onto subfunctions

* More generic approach using fieldsplit

* Interpolate: Fix parloop args ordering

* Fix dual

* Interpolate from MixedFunctionSpace

* Test Interpolate from vector/mixed to mixed

* Interpolate from (scalar/mixed) to (scalar/mixed)

Co-authored-by: Pablo Brubeck <brubeck@protonmail.com>

---------

Co-authored-by: David A. Ham <david.ham@imperial.ac.uk>

Author: pbrubeck

firedrake/adjoint_utils/blocks/assembly.py

@@ -101,8 +101,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         c = block_variable.output
         c_rep = block_variable.saved_output
 
-        from ufl.algorithms.analysis import extract_arguments
-        arity_form = len(extract_arguments(form))
+        arity_form = len(form.arguments())
 
         if isconstant(c):
             mesh = as_domain(self.form)
@@ -157,8 +156,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
         hessian_input = hessian_inputs[0]
         adj_input = adj_inputs[0]
 
-        from ufl.algorithms.analysis import extract_arguments
-        arity_form = len(extract_arguments(form))
+        arity_form = len(form.arguments())
 
         c1 = block_variable.output
         c1_rep = block_variable.saved_output

firedrake/assemble.py

@@ -18,6 +18,7 @@
 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
@@ -230,34 +231,36 @@ def assemble(self, tensor=None, current_state=None):
                 # Only Expr resulting in a Matrix if assembled are BaseFormOperator
                 if not all(isinstance(op, matrix.AssembledMatrix) for op in (a, b)):
                     raise TypeError('Mismatching Sum shapes')
-                return get_assembler(ufl.FormSum((a, 1), (b, 1))).assemble()
+                return assemble(ufl.FormSum((a, 1), (b, 1)), tensor=tensor)
             elif isinstance(expr, ufl.algebra.Product):
                 a, b = expr.ufl_operands
                 scalar = [e for e in expr.ufl_operands if is_scalar_constant_expression(e)]
                 if scalar:
                     base_form = a if a is scalar else b
                     assembled_mat = assemble(base_form)
-                    return get_assembler(ufl.FormSum((assembled_mat, scalar[0]))).assemble()
+                    return assemble(ufl.FormSum((assembled_mat, scalar[0])), tensor=tensor)
                 a, b = [assemble(e) for e in (a, b)]
-                return get_assembler(ufl.action(a, b)).assemble()
+                return assemble(ufl.action(a, b), tensor=tensor)
         # -- Linear combination of Functions and 1-form BaseFormOperators -- #
         # Example: a * u1 + b * u2 + c * N(u1; v*) + d * N(u2; v*)
         # with u1, u2 Functions, N a BaseFormOperator, and a, b, c, d scalars or 0-form BaseFormOperators.
         else:
             base_form_operators = extract_base_form_operators(expr)
-            assembled_bfops = [firedrake.assemble(e) for e in base_form_operators]
             # Substitute base form operators with their output before examining the expression
             # which avoids conflict when determining function space, for example:
             # extract_coefficients(Interpolate(u, V2)) with u \in V1 will result in an output function space V1
             # instead of V2.
             if base_form_operators:
-                expr = ufl.replace(expr, dict(zip(base_form_operators, assembled_bfops)))
-            try:
-                coefficients = ufl.algorithms.extract_coefficients(expr)
-                V, = set(c.function_space() for c in coefficients) - {None}
-            except ValueError:
-                raise ValueError("Cannot deduce correct target space from pointwise expression")
-            return firedrake.Function(V).assign(expr)
+                assembled_bfops = {e: firedrake.assemble(e) for e in base_form_operators}
+                expr = ufl.replace(expr, assembled_bfops)
+            if tensor is None:
+                try:
+                    coefficients = ufl.algorithms.extract_coefficients(expr)
+                    V, = set(c.function_space() for c in coefficients) - {None}
+                except ValueError:
+                    raise ValueError("Cannot deduce correct target space from pointwise expression")
+                tensor = firedrake.Function(V)
+            return tensor.assign(expr)
 
 
 class AbstractFormAssembler(abc.ABC):
@@ -493,9 +496,10 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             if all(isinstance(op, numbers.Complex) for op in args):
                 result = sum(weight * arg for weight, arg in zip(expr.weights(), args))
                 return tensor.assign(result) if tensor else result
-            elif all(isinstance(op, firedrake.Cofunction) for op in args):
+            elif (all(isinstance(op, firedrake.Cofunction) for op in args)
+                    or all(isinstance(op, firedrake.Function) for op in args)):
                 V, = set(a.function_space() for a in args)
-                result = tensor if tensor else firedrake.Cofunction(V)
+                result = tensor if tensor else firedrake.Function(V)
                 result.dat.maxpy(expr.weights(), [a.dat for a in args])
                 return result
             elif all(isinstance(op, ufl.Matrix) for op in args):
@@ -540,7 +544,10 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             if assembled_expression:
                 # Occur in situations such as Interpolate composition
                 expression = assembled_expression[0]
-            expr = expr._ufl_expr_reconstruct_(expression, v)
+
+            reconstruct_interp = expr._ufl_expr_reconstruct_
+            if (v, expression) != expr.argument_slots():
+                expr = reconstruct_interp(expression, v=v)
 
             # Different assembly procedures:
             # 1) Interpolate(Argument(V1, 1), Argument(V2.dual(), 0)) -> Jacobian (Interpolate matrix)
@@ -552,27 +559,59 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             # If argument numbers have been swapped => Adjoint.
             arg_expression = ufl.algorithms.extract_arguments(expression)
             is_adjoint = (arg_expression and arg_expression[0].number() == 0)
+
+            # Dual interpolation from mixed source
+            if is_adjoint and len(expr.function_space()) > 1:
+                cur = 0
+                sub_expressions = []
+                components = numpy.reshape(expression, (-1,))
+                for Vi in expr.function_space():
+                    sub_expressions.append(ufl.as_tensor(components[cur:cur+Vi.value_size].reshape(Vi.value_shape)))
+                    cur += Vi.value_size
+
+                # Component-split of the primal expression interpolated into the dual argument-split
+                split_interp = sum(reconstruct_interp(sub_expressions[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_expression[0].function_space()) > 1 and rank == 1:
+                V = arg_expression[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 primal space
+            V = expr.function_space()
+            if is_adjoint or rank == 0:
+                V = V.dual()
+
             # Workaround: Renumber argument when needed since Interpolator assumes it takes a zero-numbered argument.
-            if not is_adjoint and rank != 1:
+            if not is_adjoint and rank == 2:
                 _, v1 = expr.arguments()
-                expression = ufl.replace(expression, {v1: firedrake.Argument(v1.function_space(), number=0, part=v1.part())})
+                expression = ufl.replace(expression, {v1: v1.reconstruct(number=0)})
             # Get the interpolator
             interp_data = expr.interp_data
             default_missing_val = interp_data.pop('default_missing_val', None)
-            interpolator = firedrake.Interpolator(expression, expr.function_space(), **interp_data)
+            interpolator = firedrake.Interpolator(expression, V, **interp_data)
             # Assembly
-            if rank == 1:
+            if rank == 0:
+                Iu = interpolator._interpolate(default_missing_val=default_missing_val)
+                return assemble(ufl.Action(v, Iu), tensor=tensor)
+            elif rank == 1:
                 # Assembling the action of the Jacobian adjoint.
                 if is_adjoint:
-                    output = tensor or firedrake.Cofunction(arg_expression[0].function_space().dual())
-                    return interpolator._interpolate(v, output=output, adjoint=True, default_missing_val=default_missing_val)
+                    return interpolator._interpolate(v, output=tensor, adjoint=True, default_missing_val=default_missing_val)
                 # Assembling the Jacobian action.
-                if interpolator.nargs:
+                elif interpolator.nargs:
                     return interpolator._interpolate(expression, output=tensor, default_missing_val=default_missing_val)
                 # Assembling the operator
-                if tensor is None:
+                elif tensor is None:
                     return interpolator._interpolate(default_missing_val=default_missing_val)
-                return firedrake.Interpolator(expression, tensor, **interp_data)._interpolate(default_missing_val=default_missing_val)
+                else:
+                    return firedrake.Interpolator(expression, tensor, **interp_data)._interpolate(default_missing_val=default_missing_val)
             elif rank == 2:
                 res = tensor.petscmat if tensor else PETSc.Mat()
                 # Get the interpolation matrix
@@ -799,7 +838,7 @@ def restructure_base_form(expr, visited=None):
                 replace_map = {arg: left}
                 # Decrease number for all the other arguments since the lowest numbered argument will be replaced.
                 other_args = [a for a in right.arguments() if a is not arg]
-                new_args = [firedrake.Argument(a.function_space(), number=a.number()-1, part=a.part()) for a in other_args]
+                new_args = [a.reconstruct(number=a.number()-1) for a in other_args]
                 replace_map.update(dict(zip(other_args, new_args)))
                 # Replace arguments
                 return ufl.replace(right, replace_map)
@@ -810,13 +849,13 @@ def restructure_base_form(expr, visited=None):
             u, v = B.arguments()
             # Let V1 and V2 be primal spaces, B: V1 -> V2 and B*: V2* -> V1*:
             # Adjoint(B(Argument(V1, 1), Argument(V2.dual(), 0))) = B(Argument(V1, 0), Argument(V2.dual(), 1))
-            reordered_arguments = (firedrake.Argument(u.function_space(), number=v.number(), part=v.part()),
-                                   firedrake.Argument(v.function_space(), number=u.number(), part=u.part()))
+            reordered_arguments = {u: u.reconstruct(number=v.number()),
+                                   v: v.reconstruct(number=u.number())}
             # Replace arguments in argument slots
-            return ufl.replace(B, dict(zip((u, v), reordered_arguments)))
+            return ufl.replace(B, reordered_arguments)
 
         # -- Case (5) -- #
-        if isinstance(expr, ufl.core.base_form_operator.BaseFormOperator) and not expr.arguments():
+        if isinstance(expr, ufl.core.base_form_operator.BaseFormOperator) and len(expr.arguments()) == 0:
             # We are assembling a BaseFormOperator of rank 0 (no arguments).
             # B(f, u*) be a BaseFormOperator with u* a Cofunction and f a Coefficient, then:
             #    B(f, u*) <=> Action(B(f, v*), f) where v* is a Coargument

firedrake/interpolation.py

@@ -85,14 +85,11 @@ def __init__(self, expr, v,
                               the :meth:`interpolate` method or (b) set to zero.
                               Ignored if interpolating within the same mesh or onto a :func:`.VertexOnlyMesh`.
         """
-
         # Check function space
         if isinstance(v, functionspaceimpl.WithGeometry):
-            v = Argument(v.dual(), 0)
-
-        # Get the primal space (V** = V)
-        vv = v if not isinstance(v, ufl.Form) else v.arguments()[0]
-        self._function_space = vv.function_space().dual()
+            expr_args = extract_arguments(ufl.as_ufl(expr))
+            is_adjoint = len(expr_args) and expr_args[0].number() == 0
+            v = Argument(v.dual(), 1 if is_adjoint else 0)
         super().__init__(expr, v)
 
         # -- Interpolate data (e.g. `subset` or `access`) -- #
@@ -101,8 +98,7 @@ def __init__(self, expr, v,
                             "allow_missing_dofs": allow_missing_dofs,
                             "default_missing_val": default_missing_val}
 
-    def function_space(self):
-        return self._function_space
+    function_space = ufl.Interpolate.ufl_function_space
 
     def _ufl_expr_reconstruct_(self, expr, v=None, **interp_data):
         interp_data = interp_data or self.interp_data.copy()
@@ -293,9 +289,7 @@ def __init__(
         expr_args = extract_arguments(expr)
         if expr_args and expr_args[0].number() == 0:
             v, = expr_args
-            expr = replace(expr, {v: Argument(v.function_space(),
-                                              number=1,
-                                              part=v.part())})
+            expr = replace(expr, {v: v.reconstruct(number=1)})
         self.expr_renumbered = expr
 
     def _interpolate_future(self, *function, transpose=None, adjoint=False, default_missing_val=None):
@@ -870,7 +864,7 @@ def _interpolate(self, *function, output=None, transpose=None, adjoint=False, **
                 raise ValueError("The expression had arguments: we therefore need to be given a Function (not an expression) to interpolate!")
             if adjoint:
                 mul = assembled_interpolator.handle.multHermitian
-                V = self.arguments[0].function_space()
+                V = self.arguments[0].function_space().dual()
             else:
                 mul = assembled_interpolator.handle.mult
                 V = self.V
@@ -1123,34 +1117,14 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     name = kernel.name
     kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True,
                         flop_count=kernel.flop_count, events=(kernel.event,))
+
     parloop_args = [kernel, cell_set]
 
     coefficients = tsfc_interface.extract_numbered_coefficients(expr, coefficient_numbers)
     if needs_external_coords:
         coefficients = [source_mesh.coordinates] + coefficients
 
-    if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
-        if target_mesh is not source_mesh:
-            # NOTE: TSFC will sometimes drop run-time arguments in generated
-            # kernels if they are deemed not-necessary.
-            # FIXME: Checking for argument name in the inner kernel to decide
-            # whether to add an extra coefficient is a stopgap until
-            # compile_expression_dual_evaluation
-            #   (a) outputs a coefficient map to indicate argument ordering in
-            #       parloops as `compile_form` does and
-            #   (b) allows the dual evaluation related coefficients to be supplied to
-            #       them rather than having to be added post-hoc (likely by
-            #       replacing `to_element` with a CoFunction/CoArgument as the
-            #       target `dual` which would contain `dual` related
-            #       coefficient(s))
-            if rt_var_name in [arg.name for arg in kernel.code[name].args]:
-                # Add the coordinates of the target mesh quadrature points in the
-                # source mesh's reference cell as an extra argument for the inner
-                # loop. (With a vertex only mesh this is a single point for each
-                # vertex cell.)
-                coefficients.append(target_mesh.reference_coordinates)
-
-    if tensor in set((c.dat for c in coefficients)):
+    if any(c.dat == tensor for c in coefficients):
         output = tensor
         tensor = op2.Dat(tensor.dataset)
         if access is not op2.WRITE:
@@ -1196,6 +1170,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     if needs_cell_sizes:
         cs = source_mesh.cell_sizes
         parloop_args.append(cs.dat(op2.READ, cs.cell_node_map()))
+
     for coefficient in coefficients:
         if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
             coeff_mesh = extract_unique_domain(coefficient)
@@ -1227,6 +1202,30 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     for const in extract_firedrake_constants(expr):
         parloop_args.append(const.dat(op2.READ))
 
+    # Finally, add the target mesh reference coordinates if they appear in the kernel
+    if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
+        if target_mesh is not source_mesh:
+            # NOTE: TSFC will sometimes drop run-time arguments in generated
+            # kernels if they are deemed not-necessary.
+            # FIXME: Checking for argument name in the inner kernel to decide
+            # whether to add an extra coefficient is a stopgap until
+            # compile_expression_dual_evaluation
+            #   (a) outputs a coefficient map to indicate argument ordering in
+            #       parloops as `compile_form` does and
+            #   (b) allows the dual evaluation related coefficients to be supplied to
+            #       them rather than having to be added post-hoc (likely by
+            #       replacing `to_element` with a CoFunction/CoArgument as the
+            #       target `dual` which would contain `dual` related
+            #       coefficient(s))
+            if any(arg.name == rt_var_name for arg in kernel.code[name].args):
+                # Add the coordinates of the target mesh quadrature points in the
+                # source mesh's reference cell as an extra argument for the inner
+                # loop. (With a vertex only mesh this is a single point for each
+                # vertex cell.)
+                target_ref_coords = target_mesh.reference_coordinates
+                m_ = target_ref_coords.cell_node_map()
+                parloop_args.append(target_ref_coords.dat(op2.READ, m_))
+
     parloop = op2.ParLoop(*parloop_args)
     parloop_compute_callable = parloop.compute
     if isinstance(tensor, op2.Mat):

firedrake/ufl_expr.py

@@ -336,13 +336,9 @@ def adjoint(form, reordered_arguments=None, derivatives_expanded=None):
         # given.  To avoid that, always pass reordered_arguments with
         # firedrake.Argument objects.
         if reordered_arguments is None:
-            v, u = extract_arguments(form)
-            reordered_arguments = (Argument(u.function_space(),
-                                            number=v.number(),
-                                            part=v.part()),
-                                   Argument(v.function_space(),
-                                            number=u.number(),
-                                            part=u.part()))
+            v, u = form.arguments()
+            reordered_arguments = (u.reconstruct(number=v.number()),
+                                   v.reconstruct(number=u.number()))
         return ufl.adjoint(form, reordered_arguments, derivatives_expanded=derivatives_expanded)