Fix assemble(slate.Tensor, diagonal=True)

Author: pbrubeck

firedrake/assemble.py

@@ -1183,6 +1183,8 @@ def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing
     @FormAssembler._skip_if_initialised
     def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True,
                  zero_bc_nodes=True, diagonal=False, weight=1.0):
+        if diagonal and isinstance(form, slate.TensorBase) and len(form.arguments()) == 2:
+            form = slate.DiagonalTensor(form)
         super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing)
         self._weight = weight
         self._diagonal = diagonal

firedrake/slate/slate.py

@@ -1387,13 +1387,12 @@ def arg_function_spaces(self):
         """Returns a tuple of function spaces that the tensor
         is defined on.
         """
-        tensor, = self.operands
-        return tuple(arg.function_space() for arg in tensor.arguments())
+        return tuple(arg.function_space() for arg in self.arguments())
 
     def arguments(self):
         """Returns a tuple of arguments associated with the tensor."""
         tensor, = self.operands
-        return tensor.arguments()
+        return tensor.arguments()[:1]
 
     def _output_string(self, prec=None):
         """Creates a string representation of the diagonal of a tensor."""

tests/firedrake/slate/test_linear_algebra.py

@@ -152,10 +152,8 @@ def test_inverse_action(mat_type, rhs_type):
     assert np.allclose(x.dat.data, f.dat.data, rtol=1.e-13)
 
 
-@pytest.mark.parametrize("mat_type, rhs_type", [
-    ("slate", "slate"), ("slate", "form"), ("slate", "cofunction"),
-    ("aij", "cofunction"), ("aij", "form"),
-    ("matfree", "cofunction"), ("matfree", "form")])
+@pytest.mark.parametrize("rhs_type", ["slate", "form", "cofunction"])
+@pytest.mark.parametrize("mat_type", ["slate", "aij", "matfree"])
 def test_solve_interface(mat_type, rhs_type):
     mesh = UnitSquareMesh(1, 1)
     V = FunctionSpace(mesh, "HDivT", 0)
@@ -180,12 +178,8 @@ def test_solve_interface(mat_type, rhs_type):
     else:
         raise ValueError("Invalid rhs type")
 
-    sp = None
-    if mat_type == "matfree":
-        sp = {"pc_type": "none"}
-
     x = Function(V)
     problem = LinearVariationalProblem(A, b, x, bcs=bcs)
-    solver = LinearVariationalSolver(problem, solver_parameters=sp)
+    solver = LinearVariationalSolver(problem)
     solver.solve()
     assert np.allclose(x.dat.data, f.dat.data, rtol=1.e-13)