From f2305dc180a1f4453ac4b27467796217cbf97371 Mon Sep 17 00:00:00 2001
From: Josh Hope-Collins <joshua.hope-collins13@imperial.ac.uk>
Date: Tue, 17 Dec 2024 07:33:51 +0000
Subject: [PATCH] test matrix-free fieldsplit with cofunction rhs

---
 .../regression/test_fieldsplit_cofunction.py  | 60 -------------------
 .../firedrake/regression/test_matrix_free.py  | 10 +++-
 2 files changed, 8 insertions(+), 62 deletions(-)
 delete mode 100644 tests/firedrake/regression/test_fieldsplit_cofunction.py

diff --git a/tests/firedrake/regression/test_fieldsplit_cofunction.py b/tests/firedrake/regression/test_fieldsplit_cofunction.py
deleted file mode 100644
index 4f61fb24d4..0000000000
--- a/tests/firedrake/regression/test_fieldsplit_cofunction.py
+++ /dev/null
@@ -1,60 +0,0 @@
-import firedrake as fd
-
-
-def test_fieldsplit_cofunction():
-    """
-    Test that fieldsplit preconditioners can be used
-    with a cofunction on the right hand side.
-    """
-    mesh = fd.UnitSquareMesh(4, 4)
-    BDM = fd.FunctionSpace(mesh, "BDM", 1)
-    DG = fd.FunctionSpace(mesh, "DG", 0)
-    W = BDM*DG
-
-    u, p = fd.TrialFunctions(W)
-    v, q = fd.TestFunctions(W)
-
-    # simple wave equation scheme
-    a = (fd.dot(u, v) + fd.div(v)*p
-         - fd.div(u)*q + p*q)*fd.dx
-
-    x, y = fd.SpatialCoordinate(mesh)
-
-    f = fd.Function(W)
-
-    f.subfunctions[0].project(
-        fd.as_vector([0.01*y, 0]))
-    f.subfunctions[1].interpolate(
-        -10*fd.exp(-(pow(x - 0.5, 2) + pow(y - 0.5, 2)) / 0.02))
-
-    # compare to plain 1-form
-    L_check = fd.inner(f, fd.TestFunction(W))*fd.dx
-    L_cofun = f.riesz_representation()
-
-    # brute force schur complement solver
-    params = {
-        'ksp_converged_reason': None,
-        'ksp_type': 'preonly',
-        'pc_type': 'fieldsplit',
-        'pc_fieldsplit_type': 'schur',
-        'pc_fieldsplit_schur_fact_type': 'full',
-        'pc_fieldsplit_schur_precondition': 'full',
-        'fieldsplit': {
-            'ksp_type': 'preonly',
-            'pc_type': 'lu'
-        }
-    }
-
-    w_check = fd.Function(W)
-    problem_check = fd.LinearVariationalProblem(a, L_check, w_check)
-    solver_check = fd.LinearVariationalSolver(problem_check,
-                                              solver_parameters=params)
-    solver_check.solve()
-
-    w_cofun = fd.Function(W)
-    problem_cofun = fd.LinearVariationalProblem(a, L_cofun, w_cofun)
-    solver_cofun = fd.LinearVariationalSolver(problem_cofun,
-                                              solver_parameters=params)
-    solver_cofun.solve()
-
-    assert fd.errornorm(w_check, w_cofun) < 1e-14
diff --git a/tests/firedrake/regression/test_matrix_free.py b/tests/firedrake/regression/test_matrix_free.py
index 6cc6b207f7..486b337fb2 100644
--- a/tests/firedrake/regression/test_matrix_free.py
+++ b/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)