add a test

Author: pbrubeck

firedrake/slate/slate.py

@@ -1314,12 +1314,12 @@ def __new__(cls, A, B, decomposition=None):
         if A.shape[1] != B.shape[0]:
             raise ValueError(f"Illegal op on a {A.shape}-tensor with a {B.shape}-tensor.")
 
-        fsA = A.arg_function_spaces[1]
+        fsA = A.arg_function_spaces[0]
         fsB = B.arg_function_spaces[0]
 
-        assert fsA.dual() == fsB, (
+        assert fsA == fsB, (
             "Cannot perform argument contraction over middle indices. "
-            "They should be in dual function spaces."
+            "They must be in the same function space."
         )
 
         # For matrices smaller than 5x5, exact formulae can be used
@@ -1341,7 +1341,7 @@ def __init__(self, A, B, decomposition=None):
 
         super(Solve, self).__init__(A_factored, B)
 
-        Ainv_args = [a.reconstruct(a.function_space().dual()) for a in reversed(A.arguments())]
+        Ainv_args = tuple(a.reconstruct(a.function_space().dual()) for a in reversed(A.arguments()))
         self._args = Ainv_args[:-1] + B.arguments()[1:]
         self._arg_fs = [arg.function_space() for arg in self._args]
 

tests/firedrake/slate/test_assemble_tensors.py

@@ -130,6 +130,24 @@ def test_assemble_matrix(rank_two_tensor):
     assert np.allclose(M.M.values, assemble(rank_two_tensor.form).M.values, rtol=1e-14)
 
 
+def test_assemble_solve(mesh):
+    V = FunctionSpace(mesh, "DG", 0)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+
+    M = inner(u, v)*dx
+    f = Cofunction(V.dual())
+    f.dat.data[...] = 1
+
+    u1 = Function(V)
+    u2 = Function(V)
+    # Assemble a SLATE tensor into f
+    assemble(Inverse(Tensor(M)) * AssembledVector(f), tensor=u1)
+    # Assemble a different tensor into f
+    solve(M == f, u2)
+    assert np.allclose(u1.dat.data, u2.dat.data, rtol=1e-14)
+
+
 def test_assemble_vector_into_tensor(mesh):
     V = FunctionSpace(mesh, "DG", 1)
     v = TestFunction(V)