Interpolate: Fix parloop args ordering

Author: pbrubeck

firedrake/assemble.py

@@ -494,9 +494,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):
@@ -541,7 +542,8 @@ 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)
+            if (v, expression) != expr.argument_slots():
+                expr = expr._ufl_expr_reconstruct_(expression, v=v)
 
             # Different assembly procedures:
             # 1) Interpolate(Argument(V1, 1), Argument(V2.dual(), 0)) -> Jacobian (Interpolate matrix)
@@ -553,8 +555,7 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             arg_expression = ufl.algorithms.extract_arguments(expression)
 
             # Handle interpolation onto subfunctions
-            if (arg_expression and len(arg_expression[0].function_space()) > 1
-                    and not isinstance(expression, firedrake.Argument)):
+            if arg_expression and len(arg_expression[0].function_space()) > 1:
                 assert rank == 1
                 V = arg_expression[0].function_space()
                 if tensor is not None:

firedrake/interpolation.py

@@ -85,14 +85,9 @@ 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()
         super().__init__(expr, v)
 
         # -- Interpolate data (e.g. `subset` or `access`) -- #
@@ -101,8 +96,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()
@@ -1123,12 +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
 
+    needs_target_ref_coords = False
     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
@@ -1143,14 +1139,10 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
             #       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(arg.name == rt_var_name for arg in kernel.code[name].args):
+                needs_target_ref_coords = True
+
+    if tensor in set(c.dat for c in coefficients):
         output = tensor
         tensor = op2.Dat(tensor.dataset)
         if access is not op2.WRITE:
@@ -1196,6 +1188,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     if needs_cell_sizes:
         cs = target_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 +1220,15 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     for const in extract_firedrake_constants(expr):
         parloop_args.append(const.dat(op2.READ))
 
+    if needs_target_ref_coords:
+        # 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)
 
 

firedrake/variational_solver.py

@@ -9,11 +9,11 @@
     DEFAULT_SNES_PARAMETERS
 )
 from firedrake.function import Function
+from firedrake.interpolation import Interpolate
 from firedrake.matrix import MatrixBase
-from firedrake.ufl_expr import TrialFunction, TestFunction, action
+from firedrake.ufl_expr import TrialFunction, TestFunction
 from firedrake.bcs import DirichletBC, EquationBC, extract_subdomain_ids, restricted_function_space
 from firedrake.adjoint_utils import NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin
-from firedrake.__future__ import interpolate
 from ufl import replace, Form
 
 __all__ = ["LinearVariationalProblem",
@@ -100,7 +100,7 @@ def __init__(self, F, u, bcs=None, J=None,
                 F_arg, = F.arguments()
                 self.F = replace(F, {F_arg: v_res, self.u: self.u_restrict})
             else:
-                self.F = action(replace(F, {self.u: self.u_restrict}), interpolate(v_res, V))
+                self.F = Interpolate(v_res, replace(F, {self.u: self.u_restrict}))
             v_arg, u_arg = self.J.arguments()
             self.J = replace(self.J, {v_arg: v_res, u_arg: u_res, self.u: self.u_restrict})
             if self.Jp:

tests/firedrake/regression/test_interp_dual.py

@@ -280,7 +280,7 @@ def test_solve_interp_u(mesh):
 
 
 @pytest.mark.parametrize("family0,degree0,family1,degree1",
-                         [("DG", 1, "CG", 1),
+                         [("DG", 1, "CG", 2),
                           ("DG", 0, "RT", 1)])
 def test_interp_subfunction(mesh, family0, degree0, family1, degree1):
     V = FunctionSpace(mesh, "DG", 0)
@@ -291,20 +291,20 @@ def test_interp_subfunction(mesh, family0, degree0, family1, degree1):
     W1 = FunctionSpace(mesh, family1, degree1)
     W = W0 * W1
     w = TestFunction(W)
-    c = Cofunction(W.dual())
 
     expr = sum(w[i] for i in np.ndindex(w.ufl_shape))
 
-    Fw = inner(1, expr)*dx
+    Fw = inner(1, expr)*dx(degree=0)
     expected = assemble(Fw)
 
     IFv = Interpolate(expr, Fv)
 
+    c = Cofunction(W.dual())
+    c.assign(99)
     for tensor in (None, c):
-        if tensor:
-            tensor.assign(99)
-
         result = assemble(IFv, tensor=tensor)
         assert result.function_space() == W.dual()
-        with result.dat.vec_ro as x, expected.dat.vec_ro as y:
-            assert np.allclose(x[:], y[:])
+        if tensor:
+            assert result is tensor
+        for x, y, in zip(result.subfunctions, expected.subfunctions):
+            assert np.allclose(x.dat.data_ro, y.dat.data_ro)