SLATE: prevent Cofunction reassembly

Author: pbrubeck

firedrake/adjoint_utils/variational_solver.py

@@ -27,7 +27,7 @@ def wrapper(self, *args, **kwargs):
                     # Try again without expanding derivatives,
                     # as dFdu might have been simplied to an empty Form
                     self._ad_adj_F = adjoint(dFdu, derivatives_expanded=True)
-            except (TypeError, NotImplementedError):
+            except (ValueError, TypeError, NotImplementedError):
                 self._ad_adj_F = None
             self._ad_kwargs = {'Jp': self.Jp, 'form_compiler_parameters': self.form_compiler_parameters, 'is_linear': self.is_linear}
             self._ad_count_map = {}

firedrake/assemble.py

@@ -838,7 +838,11 @@ def preprocess_base_form(expr, mat_type=None, form_compiler_parameters=None):
         if mat_type != "matfree":
             # Don't expand derivatives if `mat_type` is 'matfree'
             # For "matfree", Form evaluation is delayed
-            expr = BaseFormAssembler.expand_derivatives_form(expr, form_compiler_parameters)
+            try:
+                expr = BaseFormAssembler.expand_derivatives_form(expr, form_compiler_parameters)
+            except ValueError:
+                # BaseForms with SLATE tensors are not fully supported in UFL.
+                pass
         if not isinstance(expr, (ufl.form.Form, slate.TensorBase)):
             # => No restructuring needed for Form and slate.TensorBase
             expr = BaseFormAssembler.restructure_base_form_preorder(expr)

firedrake/slate/slate.py

@@ -1430,7 +1430,9 @@ def as_slate(F):
         return F
     elif isinstance(F, Form):
         return Tensor(F)
-    elif isinstance(F, (Function, Cofunction)):
+    elif isinstance(F, Function):
+        # Do not implicitly cast Cofunctions as Slate tensors
+        # because assemble(AssembledVector(F)) repeats element summation on F
         return AssembledVector(F)
     else:
         raise TypeError(f"Cannot convert {type(F).__name__} into a slate.Tensor")

firedrake/variational_solver.py

@@ -132,14 +132,22 @@ def dm(self):
         return self.u_restrict.function_space().dm
 
     @staticmethod
-    def compute_bc_lifting(J, u):
+    def compute_bc_lifting(J, u, L=0):
         """Return the action of the bilinear form J (without bcs) on a Function u."""
         if isinstance(J, MatrixBase) and J.has_bcs:
             # Extract the full form without bcs
             if not isinstance(J.a, (ufl.BaseForm, slate.slate.TensorBase)):
                 raise TypeError(f"Could not remove bcs from {type(J).__name__}.")
             J = J.a
-        return ufl_expr.action(J, u)
+        F = ufl_expr.action(J, u)
+        if L != 0:
+            try:
+                F = F - L
+            except TypeError:
+                # Slate expressions do not combine with Cofunctions
+                # because assemble(AssembledVector(L)) repeats element summation on L
+                F = ufl.FormSum((F, 1), (L, -1))
+        return F
 
 
 class NonlinearVariationalSolver(OptionsManager, NonlinearVariationalSolverMixin):
@@ -399,7 +407,7 @@ def __init__(self, a, L, u, bcs=None, aP=None,
                 raise TypeError("Provided RHS is a '%s', not a Form or Slate Tensor" % type(L).__name__)
             if len(L.arguments()) != 1 and not L.empty():
                 raise ValueError("Provided RHS is not a linear form")
-            F = self.compute_bc_lifting(a, u) - L
+            F = self.compute_bc_lifting(a, u, L=L)
 
         super(LinearVariationalProblem, self).__init__(F, u, bcs=bcs, J=a, Jp=aP,
                                                        form_compiler_parameters=form_compiler_parameters,

tests/firedrake/slate/test_linear_algebra.py

@@ -150,3 +150,42 @@ def test_inverse_action(mat_type, rhs_type):
     x = Function(V)
     assemble(action(Ainv, b), tensor=x)
     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")])
+def test_solve_interface(mat_type, rhs_type):
+    mesh = UnitSquareMesh(1, 1)
+    V = FunctionSpace(mesh, "HDivT", 0)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    f = Function(V).assign(1.0)
+    bcs = DirichletBC(V, f, "on_boundary")
+
+    a = avg(inner(u, v))*dS + inner(u, v)*ds
+    A = Tensor(a)
+    if mat_type != "slate":
+        A = assemble(A, bcs=bcs, mat_type=mat_type)
+        bcs = None
+
+    L = action(a, f)
+    if rhs_type == "form":
+        b = L
+    elif rhs_type == "cofunction":
+        b = assemble(L)
+    elif rhs_type == "slate":
+        b = Tensor(L)
+    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.solve()
+    assert np.allclose(x.dat.data, f.dat.data, rtol=1.e-13)