Support tensor kwarg

Author: pbrubeck

firedrake/assemble.py

@@ -563,8 +563,11 @@ 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.
             is_adjoint = (arg_expression and arg_expression[0].number() == 0)

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)