Fix TransferManager reuse for mixed coefficients (#4600)

* Fix TransferManager reuse for mixed coefficients

Author: pbrubeck

firedrake/variational_solver.py

@@ -332,12 +332,12 @@ def solve(self, bounds=None):
         problem = self._problem
         forms = (problem.F, problem.J, problem.Jp)
         coefficients = utils.unique(chain.from_iterable(form.coefficients() for form in forms if form is not None))
-        spaces = chain.from_iterable(c.function_space() for c in coefficients)
+        coefficients += problem.u.subfunctions
         solution_dm = self.snes.getDM()
         # Grab the unique DMs for this problem
         problem_dms = []
-        for V in spaces:
-            dm = V.dm
+        for c in coefficients:
+            dm = c.function_space().dm
             if dm == solution_dm:
                 # Make sure the solution dm is visited last
                 continue

tests/firedrake/multigrid/test_transfer_manager.py

@@ -134,13 +134,20 @@ def test_transfer_manager_dat_version_cache(action, transfer_op, spaces):
         raise ValueError(f"Unrecognized action {action}")
 
 
-@pytest.mark.parametrize("family, degree", [("CG", 1), ("R", 0)])
-def test_cached_transfer(family, degree):
+@pytest.mark.parametrize("family, degree, coefficient", [
+    ("CG", 1, "repeated"),
+    ("CG", 1, "mixed"),
+    ("CG", 1, "bcs"),
+    ("R", 0, "repeated"),
+])
+def test_cached_transfer(family, degree, coefficient):
     # Test that we can properly reuse transfers within solve
-    sp = {"mat_type": "matfree",
-          "pc_type": "mg",
-          "mg_coarse_pc_type": "none",
-          "mg_levels_pc_type": "none"}
+    sp = {
+        "mat_type": "matfree",
+        "pc_type": "mg",
+        "mg_levels_pc_type": "none",
+        "mg_coarse_pc_type": "none",
+    }
 
     base = UnitSquareMesh(1, 1)
     hierarchy = MeshHierarchy(base, 3)
@@ -149,13 +156,30 @@ def test_cached_transfer(family, degree):
     V = FunctionSpace(mesh, family, degree)
     u = Function(V)
 
-    R1 = FunctionSpace(mesh, "R", 0)
-    R2 = FunctionSpace(mesh, "R", 0)
-    c1 = Function(R1).assign(1)
-    c2 = Function(R2).assign(1)
+    bcs = None
+    if coefficient == "mixed":
+        R = FunctionSpace(mesh, "R", 0)
+        R2 = R * R
+        c = Function(R2).assign(1)
+        c1 = c.subfunctions[0]
+        c2 = c[1]
+    elif coefficient == "repeated":
+        R1 = FunctionSpace(mesh, "R", 0)
+        R2 = FunctionSpace(mesh, "R", 0)
+        c1 = Function(R1).assign(1)
+        c2 = Function(R2).assign(1)
+    elif coefficient == "bcs":
+        c1 = 1
+        c2 = 1
+        R = FunctionSpace(mesh, "R", 0)
+        R2 = R * R
+        g = Function(R2).assign(1)
+        bcs = [DirichletBC(V, g[0], (1, 2)), DirichletBC(V, g[1], (3, 4))]
+    else:
+        raise ValueError(f"Unrecognized coefficient type {coefficient}")
 
     F = inner(u - 1, (c1 + c2)*TestFunction(V)) * dx
-    problem = NonlinearVariationalProblem(F, u)
+    problem = NonlinearVariationalProblem(F, u, bcs=bcs)
     solver = NonlinearVariationalSolver(problem, solver_parameters=sp)
 
     transfer = TransferManager()