Update interpolation.py

Author: pbrubeck

firedrake/assemble.py

@@ -590,28 +590,38 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             if not is_adjoint and rank == 2:
                 _, v1 = expr.arguments()
                 operand = ufl.replace(operand, {v1: v1.reconstruct(number=0)})
+
+            target_mesh = V.mesh()
+            source_mesh = extract_unique_domain(operand) or target_mesh
+            same_mesh = source_mesh.topology is target_mesh.topology
+
             # Get the interpolator
-            interp_data = expr.interp_data.copy()
+            interp_data = expr.interp_data
             default_missing_val = interp_data.pop('default_missing_val', None)
-            if (is_adjoint and rank == 1) or rank == 0:
+            if same_mesh and ((is_adjoint and rank == 1) or rank == 0):
+                interp_data = interp_data.copy()
                 interp_data["access"] = op2.INC
-            interpolator = firedrake.Interpolator(operand, v, **interp_data)
+
+            dual_arg = v if same_mesh else V
+            interp_expr = firedrake.Interpolate(operand, v=dual_arg, **interp_data)
+            interpolator = firedrake.Interpolator(interp_expr, V, **interp_data)
+
             # Assembly
             if rank == 0:
                 result = interpolator._interpolate(output=tensor, default_missing_val=default_missing_val)
                 return result.dat.data.item() if tensor is None else result
             elif rank == 1:
                 # Assembling the action of the Jacobian adjoint.
                 if is_adjoint:
-                    return interpolator._interpolate(v, output=tensor, default_missing_val=default_missing_val)
+                    return interpolator._interpolate(v, output=tensor, adjoint=True, default_missing_val=default_missing_val)
                 # Assembling the Jacobian action.
                 elif interpolator.nargs:
                     return interpolator._interpolate(operand, output=tensor, default_missing_val=default_missing_val)
                 # Assembling the operator
                 elif tensor is None:
                     return interpolator._interpolate(default_missing_val=default_missing_val)
                 else:
-                    return firedrake.Interpolator(operand, tensor, **interp_data)._interpolate(default_missing_val=default_missing_val)
+                    return firedrake.Interpolator(interp_expr, 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

firedrake/interpolation.py

@@ -241,6 +241,8 @@ class Interpolator(abc.ABC):
     """
 
     def __new__(cls, expr, V, **kwargs):
+        if isinstance(expr, ufl.Interpolate):
+            expr, = expr.ufl_operands
         target_mesh = as_domain(V)
         source_mesh = extract_unique_domain(expr) or target_mesh
         submesh_interp_implemented = \
@@ -266,6 +268,8 @@ def __init__(
         allow_missing_dofs=False,
         matfree=True
     ):
+        if isinstance(expr, ufl.Interpolate):
+            expr, = expr.ufl_operands
         self.expr = expr
         self.V = V
         self.subset = subset
@@ -374,6 +378,8 @@ def __init__(
                 "Can only interpolate into spaces with point evaluation nodes."
             )
 
+        if isinstance(expr, ufl.Interpolate):
+            expr, = expr.ufl_operands
         super().__init__(expr, V, subset, freeze_expr, access, bcs, allow_missing_dofs, matfree)
 
         self.arguments = extract_arguments(expr)
@@ -540,7 +546,7 @@ def _interpolate(
                 V_dest = self.expr.function_space().dual()
             except AttributeError:
                 if self.nargs:
-                    V_dest = self.arguments[0].function_space().dual()
+                    V_dest = self.arguments[-1].function_space().dual()
                 else:
                     coeffs = extract_coefficients(self.expr)
                     if len(coeffs):
@@ -552,8 +558,6 @@ def _interpolate(
         else:
             if isinstance(self.V, (firedrake.Function, firedrake.Cofunction)):
                 V_dest = self.V.function_space()
-            elif isinstance(self.V, firedrake.Coargument):
-                V_dest = self.V.function_space().dual()
             else:
                 V_dest = self.V
         if output:
@@ -679,10 +683,14 @@ class SameMeshInterpolator(Interpolator):
     def __init__(self, expr, V, subset=None, freeze_expr=False, access=op2.WRITE,
                  bcs=None, matfree=True, allow_missing_dofs=False, **kwargs):
         if subset is None:
+            if isinstance(expr, ufl.Interpolate):
+                operand, = expr.ufl_operands
+            else:
+                operand = expr
             target_mesh = as_domain(V)
-            source_mesh = extract_unique_domain(expr)
+            source_mesh = extract_unique_domain(operand) or target_mesh
             target = target_mesh.topology
-            source = target if source_mesh is None else source_mesh.topology
+            source = source_mesh.topology
             if all(isinstance(m, firedrake.mesh.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":
@@ -703,7 +711,7 @@ def __init__(self, expr, V, subset=None, freeze_expr=False, access=op2.WRITE,
             self.callable, arguments = make_interpolator(expr, V, subset, access, bcs=bcs, matfree=matfree)
         except FIAT.hdiv_trace.TraceError:
             raise NotImplementedError("Can't interpolate onto traces sorry")
-        self.arguments = arguments
+        self.arguments = expr.arguments()
         self.nargs = len(arguments)
 
     @PETSc.Log.EventDecorator()
@@ -735,16 +743,19 @@ def _interpolate(self, *function, output=None, transpose=None, adjoint=False, **
                     # Interpolation action
                     self.frozen_assembled_interpolator = assembled_interpolator.copy()
 
-        if self.nargs == 2:
+        if hasattr(assembled_interpolator, "handle") and len(function):
             function, = function
             if not hasattr(function, "dat"):
                 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().dual()
+                assert function.function_space() == self.arguments[1].function_space()
             else:
                 mul = assembled_interpolator.handle.mult
-                V = self.V
+                V = self.arguments[1].function_space().dual()
+                assert function.function_space() == self.arguments[0].function_space()
+
             result = output or firedrake.Function(V)
             with function.dat.vec_ro as x, result.dat.vec_wo as out:
                 if x is not out:
@@ -772,29 +783,23 @@ def _interpolate(self, *function, output=None, transpose=None, adjoint=False, **
 
 @PETSc.Log.EventDecorator()
 def make_interpolator(expr, V, subset, access, bcs=None, matfree=True):
-    assert isinstance(expr, ufl.classes.Expr)
+    assert isinstance(expr, ufl.Interpolate)
+    dual_arg, operand = expr.argument_slots()
+    assert isinstance(dual_arg, (ufl.Coargument, ufl.Cofunction))
+
+    target_mesh = as_domain(dual_arg)
+    source_mesh = extract_unique_domain(operand) or target_mesh
+    same_mesh = target_mesh is source_mesh
+
+    vom_onto_other_vom = (
+        isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology)
+        and isinstance(source_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology)
+        and target_mesh is not source_mesh
+    )
 
-    if isinstance(V, (ufl.Coargument, ufl.Cofunction)):
-        dual_arg = V
-        V = dual_arg.function_space().dual()
-    elif isinstance(V, (ufl.FunctionSpace, ufl.Coefficient)):
-        fs = V if isinstance(V, ufl.FunctionSpace) else V.function_space()
-        dual_arg = Coargument(fs.dual(), number=0)
-
-    arguments = extract_arguments(expr)
-    if isinstance(dual_arg, ufl.Coargument):
-        arguments.append(dual_arg)
+    arguments = expr.arguments()
     rank = len(arguments)
-
-    target_mesh = as_domain(V)
     if rank <= 1:
-        source_mesh = extract_unique_domain(expr) or target_mesh
-        vom_onto_other_vom = (
-            isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology)
-            and isinstance(source_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology)
-            and target_mesh is not source_mesh
-        )
-
         if rank == 0:
             R = firedrake.FunctionSpace(target_mesh, "Real", 0)
             f = firedrake.Function(R)
@@ -817,15 +822,8 @@ def make_interpolator(expr, V, subset, access, bcs=None, matfree=True):
             raise ValueError("Cannot interpolate an expression with an argument into a Function")
         if len(V) > 1:
             raise NotImplementedError("Interpolation of mixed expressions with arguments is not supported")
-
         argfs = arguments[0].function_space()
-        source_mesh = argfs.mesh()
         argfs_map = argfs.cell_node_map()
-        vom_onto_other_vom = (
-            isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology)
-            and isinstance(source_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology)
-            and target_mesh is not source_mesh
-        )
         if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology) and target_mesh is not source_mesh and not vom_onto_other_vom:
             if not isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
                 raise NotImplementedError("Can only interpolate onto a Vertex Only Mesh")
@@ -863,8 +861,10 @@ def make_interpolator(expr, V, subset, access, bcs=None, matfree=True):
     else:
         raise ValueError("Cannot interpolate an expression with %d arguments" % rank)
 
+    if not same_mesh:
+        arguments = extract_arguments(operand)
     if vom_onto_other_vom:
-        wrapper = VomOntoVomWrapper(V, source_mesh, target_mesh, expr, arguments, matfree)
+        wrapper = VomOntoVomWrapper(V, source_mesh, target_mesh, operand, arguments, matfree)
         # NOTE: get_dat_mpi_type ensures we get the correct MPI type for the
         # data, including the correct data size and dimensional information
         # (so for vector function spaces in 2 dimensions we might need a
@@ -874,13 +874,13 @@ def make_interpolator(expr, V, subset, access, bcs=None, matfree=True):
             # when it is called.
             assert f.dat is tensor
             wrapper.mpi_type, _ = get_dat_mpi_type(f.dat)
-            assert not len(arguments)
+            assert len(arguments) == 0
 
             def callable():
                 wrapper.forward_operation(f.dat)
                 return f
         else:
-            assert rank == 2
+            assert len(arguments) == 1
             assert tensor is None
             # we know we will be outputting either a function or a cofunction,
             # both of which will use a dat as a data carrier. At present, the
@@ -904,38 +904,40 @@ def callable():
         # Make sure we have an expression of the right length i.e. a value for
         # each component in the value shape of each function space
         loops = []
-        if numpy.prod(expr.ufl_shape, dtype=int) != V.value_size:
+        if numpy.prod(operand.ufl_shape, dtype=int) != V.value_size:
             raise RuntimeError('Expression of length %d required, got length %d'
-                               % (V.value_size, numpy.prod(expr.ufl_shape, dtype=int)))
+                               % (V.value_size, numpy.prod(operand.ufl_shape, dtype=int)))
 
         if len(V) == 1:
-            loops.extend(_interpolator(V, tensor, expr, dual_arg, subset, arguments, access, bcs=bcs))
+            loops.extend(_interpolator(V, tensor, expr, subset, arguments, access, bcs=bcs))
         else:
-            if (hasattr(expr, "subfunctions") and len(expr.subfunctions) == len(V)
-                    and all(sub_expr.ufl_shape == Vsub.value_shape for Vsub, sub_expr in zip(V, expr.subfunctions))):
+            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
-                expressions = expr.subfunctions
+                operands = operand.subfunctions
             else:
                 # Unflatten the expression into the shapes of the mixed components
                 offset = 0
-                expressions = []
+                operands = []
                 for Vsub in V:
                     if len(Vsub.value_shape) == 0:
-                        expressions.append(expr[offset])
+                        operands.append(operand[offset])
                     else:
-                        components = [expr[offset + j] for j in range(Vsub.value_size)]
-                        expressions.append(ufl.as_tensor(numpy.reshape(components, Vsub.value_shape)))
+                        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
 
             if isinstance(dual_arg, Cofunction):
                 duals = dual_arg.subfunctions
             elif isinstance(dual_arg, Coargument):
-                duals = [Coargument(Vsub.dual(), number=dual_arg.number()) for Vsub in V]
+                duals = [Coargument(Vsub, number=dual_arg.number()) for Vsub in dual_arg.function_space()]
             else:
                 raise ValueError("dual_arg must be a Cofunction or Coargument")
+
             # Interpolate each sub expression into each function space
-            for Vsub, sub_tensor, sub_expr, sub_dual in zip(V, tensor, expressions, duals):
-                loops.extend(_interpolator(Vsub, sub_tensor, sub_expr, sub_dual, subset, arguments, access, bcs=bcs))
+            for Vsub, sub_tensor, sub_op, sub_dual in zip(V, tensor, operands, duals):
+                sub_expr = ufl.Interpolate(sub_op, sub_dual)
+                loops.extend(_interpolator(Vsub, sub_tensor, sub_expr, subset, arguments, access, bcs=bcs))
 
         if bcs and rank == 1:
             loops.extend(partial(bc.apply, f) for bc in bcs)
@@ -949,11 +951,18 @@ def callable(loops, f):
 
 
 @utils.known_pyop2_safe
-def _interpolator(V, tensor, expr, dual_arg, subset, arguments, access, bcs=None):
+def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     try:
         expr = ufl.as_ufl(expr)
     except ValueError:
         raise ValueError("Expecting to interpolate a UFL expression")
+
+    interp_expr = expr
+    if isinstance(expr, ufl.Interpolate):
+        dual_arg, expr = expr.argument_slots()
+    else:
+        dual_arg = Coargument(V.dual(), number=0)
+
     try:
         to_element = create_element(V.ufl_element())
     except KeyError:
@@ -1029,7 +1038,7 @@ def _interpolator(V, tensor, expr, dual_arg, subset, arguments, access, bcs=None
     # FIXME: for the runtime unknown point set (for cross-mesh
     # interpolation) we have to pass the finat element we construct
     # here. Ideally we would only pass the UFL element through.
-    kernel = compile_expression(cell_set.comm, expr, dual_arg, to_element, V.ufl_element(),
+    kernel = compile_expression(cell_set.comm, interp_expr, to_element, V.ufl_element(),
                                 domain=source_mesh, parameters=parameters)
     ast = kernel.ast
     oriented = kernel.oriented
@@ -1042,7 +1051,7 @@ def _interpolator(V, tensor, expr, dual_arg, subset, arguments, access, bcs=None
 
     parloop_args = [kernel, cell_set]
 
-    interp_expr = ufl.Interpolate(expr, dual_arg)
+    interp_expr = ufl.Interpolate(expr, v=dual_arg)
     coefficients = tsfc_interface.extract_numbered_coefficients(interp_expr, coefficient_numbers)
     if needs_external_coords:
         coefficients = [source_mesh.coordinates] + coefficients
@@ -1061,12 +1070,13 @@ def _interpolator(V, tensor, expr, dual_arg, subset, arguments, access, bcs=None
     if isinstance(tensor, op2.Global):
         parloop_args.append(tensor(access))
     elif isinstance(tensor, op2.Dat):
-        V_dest = arguments[0].function_space() if isinstance(dual_arg, ufl.Cofunction) else V
+        V_dest = arguments[-1].function_space() if isinstance(dual_arg, ufl.Cofunction) else V
         parloop_args.append(tensor(access, V_dest.cell_node_map()))
     else:
         assert access == op2.WRITE  # Other access descriptors not done for Matrices.
         rows_map = V.cell_node_map()
         Vcol = arguments[0].function_space()
+        assert tensor.handle.getSize() == (V.dim(), Vcol.dim())
         if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
             columns_map = Vcol.cell_node_map()
             if target_mesh is not source_mesh:
@@ -1165,9 +1175,10 @@ def _interpolator(V, tensor, expr, dual_arg, subset, arguments, access, bcs=None
                                   f"firedrake-tsfc-expression-kernel-cache-uid{os.getuid()}")
 
 
-def _compile_expression_key(comm, expr, dual_arg, to_element, ufl_element, domain, parameters) -> tuple[Hashable, ...]:
+def _compile_expression_key(comm, expr, to_element, ufl_element, domain, parameters) -> tuple[Hashable, ...]:
     """Generate a cache key suitable for :func:`tsfc.compile_expression_dual_evaluation`."""
-    return (hash_expr(expr), type(dual_arg), hash(ufl_element), utils.tuplify(parameters))
+    dual_arg, operand = expr.argument_slots()
+    return (hash_expr(operand), type(dual_arg), hash(ufl_element), utils.tuplify(parameters))
 
 
 @memory_and_disk_cache(

tests/firedrake/regression/test_interpolate.py

@@ -341,7 +341,6 @@ def test_adjoint_Pk(degree):
 
     v_adj_form = assemble(interpolate(TestFunction(Pk), v * dx))
 
-    assert v_adj.function_space() == v_adj_form.function_space()
     assert np.allclose(v_adj_form.dat.data, v_adj.dat.data)
 
 
@@ -354,7 +353,6 @@ def test_adjoint_quads():
     u_P1 = assemble(conj(TestFunction(P1)) * dx)
     v_adj = assemble(interpolate(TestFunction(P1), assemble(v * dx)))
 
-    assert v_adj.function_space() == u_P1.function_space()
     assert np.allclose(u_P1.dat.data, v_adj.dat.data)
 
 
@@ -367,7 +365,6 @@ def test_adjoint_dg():
     u_cg = assemble(conj(TestFunction(cg1)) * dx)
     v_adj = assemble(interpolate(TestFunction(cg1), assemble(v * dx)))
 
-    assert v_adj.function_space() == u_cg.function_space()
     assert np.allclose(u_cg.dat.data, v_adj.dat.data)