Fix fieldsplit with Cofunction right hand side.

firedrake/formmanipulation.py

@@ -88,6 +88,7 @@ def argument(self, o):
         from ufl import split
         from firedrake import MixedFunctionSpace, FunctionSpace
         V = o.function_space()
+
         if len(V) == 1:
             # Not on a mixed space, just return ourselves.
             return o
@@ -127,6 +128,56 @@ def argument(self, o):
                          for j in numpy.ndindex(V_is[i].value_shape)]
         return self._arg_cache.setdefault(o, as_vector(args))
 
+    def cofunction(self, o):
+        from firedrake import Cofunction, FunctionSpace, MixedFunctionSpace
+        from pyop2 import MixedDat
+        from ufl.classes import ZeroBaseForm
+
+        V = o.function_space()
+
+        # Not on a mixed space, just return ourselves.
+        if len(V) == 1:
+            return o
+
+        # We only need the test space for Cofunction
+        indices = self.blocks[0]
+        V_is = V.subfunctions
+
+        try:
+            indices = tuple(indices)
+            nidx = len(indices)
+        except TypeError:  # Only one index
+            indices = (indices, )
+            nidx = 1
+
+        # the cofunction should only be returned on the
+        # diagonal elements, so if we are off-diagonal
+        # on a two-form then we return a zero form,
+        # analogously to the off components of arguments.
+        if len(self.blocks) == 2:
+            itest, itrial = self.blocks
+            on_diag = (itest == itrial)
+        else:
+            on_diag = True
+
+        # if we are on the diagonal, then return a Cofunction
+        # in the relevant subspace that points to the data in
+        # the full space. This means that the right hand side
+        # of the fieldsplit problem will be correct.
+        if on_diag:
+            if nidx == 1:
+                i = indices[0]
+                W = V_is[i]
+                W = FunctionSpace(W.mesh(), W.ufl_element())  # primal space
+                c = Cofunction(W.dual(), val=o.subfunctions[i].dat)
+            else:
+                W = MixedFunctionSpace([V_is[i] for i in indices])  # dual space
+                c = Cofunction(W, val=MixedDat(o.dat[i] for i in indices))
+        else:
+            c = ZeroBaseForm(o.arguments())
+
+        return c
+
 
 SplitForm = collections.namedtuple("SplitForm", ["indices", "form"])
 
@@ -160,6 +211,8 @@ def split_form(form, diagonal=False):
     compiler will remove these in its more complex simplification
     stages.
     """
+    from firedrake import Cofunction
+    from ufl import FormSum, Form
     splitter = ExtractSubBlock()
     args = form.arguments()
     shape = tuple(len(a.function_space()) for a in args)
@@ -168,7 +221,20 @@ def split_form(form, diagonal=False):
         assert len(shape) == 2
     for idx in numpy.ndindex(shape):
         f = splitter.split(form, idx)
-        if len(f.integrals()) > 0:
+
+        # does f actually contain anything?
+        if isinstance(f, Cofunction):
+            flen = 1
+        elif isinstance(f, FormSum):
+            flen = len(f.components())
+        elif isinstance(f, Form):
+            flen = len(f.integrals())
+        else:  # Form
+            raise ValueError(
+                "ExtractSubBlock.split should have returned an instance of "
+                "either Form, FormSum, or Cofunction")
+
+        if flen > 0:
             if diagonal:
                 i, j = idx
                 if i != j:

tests/firedrake/regression/test_matrix_free.py

@@ -130,6 +130,7 @@ def test_matrixfree_action(a, V, bcs):
 
 @pytest.mark.parametrize("preassembled", [False, True],
                          ids=["variational", "preassembled"])
+@pytest.mark.parametrize("rhs", ["func_rhs", "cofunc_rhs"])
 @pytest.mark.parametrize("parameters",
                          [{"ksp_type": "preonly",
                            "pc_type": "python",
@@ -168,7 +169,7 @@ def test_matrixfree_action(a, V, bcs):
                            "fieldsplit_1_fieldsplit_1_pc_type": "python",
                            "fieldsplit_1_fieldsplit_1_pc_python_type": "firedrake.AssembledPC",
                            "fieldsplit_1_fieldsplit_1_assembled_pc_type": "lu"}])
-def test_fieldsplitting(mesh, preassembled, parameters):
+def test_fieldsplitting(mesh, preassembled, parameters, rhs):
     V = FunctionSpace(mesh, "CG", 1)
     P = FunctionSpace(mesh, "DG", 0)
     Q = VectorFunctionSpace(mesh, "DG", 1)
@@ -184,7 +185,12 @@ def test_fieldsplitting(mesh, preassembled, parameters):
 
     a = inner(u, v)*dx
 
-    L = inner(expect, v)*dx
+    if rhs == 'func_rhs':
+        L = inner(expect, v)*dx
+    elif rhs == 'cofunc_rhs':
+        L = expect.riesz_representation()
+    else:
+        raise ValueError("Unknown right hand side type")
 
     f = Function(W)