Handle tensor kwarg

Author: pbrubeck

firedrake/assemble.py

@@ -549,9 +549,8 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             # 4) Interpolate(Argument(V1, 0), Cofunction(...)) -> Action of the Jacobian adjoint
             # This can be generalized to the case where the first slot is an arbitray expression.
             rank = len(expr.arguments())
-            arg_expression = ufl.algorithms.extract_arguments(expression)
 
-            # Handle interpolation of subfunctions
+            # Handle interpolation onto subfunctions
             parent_tensor = None
             if isinstance(expression, ufl.classes.Indexed):
                 assert rank == 1
@@ -563,10 +562,14 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 # Symbolic indirection for the input expression
                 expression = firedrake.Argument(V.sub(index), number=A.number(), part=A.part())
                 # Symbolic indirection for the output tensor
+                if tensor is not None:
+                    assert tensor.function_space() == V.dual()
+                    tensor.zero()
                 parent_tensor = tensor or firedrake.Cofunction(V.dual())
-                tensor = parent_tensor.sub(index)
+                tensor = parent_tensor.subfunctions[index]
 
             # If argument numbers have been swapped => Adjoint.
+            arg_expression = ufl.algorithms.extract_arguments(expression)
             is_adjoint = (arg_expression and arg_expression[0].number() == 0)
             # Workaround: Renumber argument when needed since Interpolator assumes it takes a zero-numbered argument.
             if not is_adjoint and rank != 1:
@@ -575,7 +578,6 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             # 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)
             # Assembly
             if rank == 1:

tests/firedrake/regression/test_interp_dual.py

@@ -279,17 +279,24 @@ def test_solve_interp_u(mesh):
     assert np.allclose(u.dat.data, u2.dat.data)
 
 
-def test_interp_subfunction(mesh):
+@pytest.mark.parametrize("tensor", (False, True), ids=("tensor=none", "tensor"))
+def test_interp_subfunction(mesh, tensor):
     V = FunctionSpace(mesh, "DG", 0)
     v = TestFunction(V)
     Fv = inner(1, v)*dx
 
-    W = V*V
+    W = V * V
     w = TestFunction(W)
-    Fw = inner(1, w[0])*dx
+    Fw = inner(1, w[1])*dx
+    expected = assemble(Fw)
 
-    I = Interpolate(w[0], Fv)
+    IFv = Interpolate(w[1], Fv)
 
-    c0 = assemble(I)
-    c1 = assemble(Fw)
-    assert np.allclose(c0.dat.data_ro, c1.dat.data_ro)
+    if tensor:
+        tensor = Cofunction(W.dual())
+    else:
+        tensor = None
+
+    c = assemble(IFv, tensor=tensor)
+    assert c.function_space() == W.dual()
+    assert np.allclose(c.dat.data_ro, expected.dat.data_ro)