Eikehmueller/gtmgpc restriction

firedrake/preconditioners/gtmg.py

@@ -107,7 +107,7 @@ def initialize(self, pc):
             interpolator = Interpolator(TestFunction(coarse_space), fine_space)
             interpolation_matrix = interpolator.callable()
             interp_petscmat = interpolation_matrix.handle
-
+        restr_petscmat = appctx.get("restriction_matrix", None)
         # We set up a PCMG object that uses the constructed interpolation
         # matrix to generate the restriction/prolongation operators.
         # This is a two-level multigrid preconditioner.
@@ -119,6 +119,8 @@ def initialize(self, pc):
         pcmg.setMGLevels(2)
         pcmg.setMGCycleType(pc.MGCycleType.V)
         pcmg.setMGInterpolation(1, interp_petscmat)
+        if restr_petscmat is not None:
+            pcmg.setMGRestriction(1, restr_petscmat)
         pcmg.setOperators(A=fine_petscmat, P=fine_petscmat)
 
         coarse_solver = pcmg.getMGCoarseSolve()

tests/firedrake/multigrid/test_poisson_gtmg.py

@@ -2,7 +2,7 @@
 import pytest
 
 
-def run_gtmg_mixed_poisson():
+def run_gtmg_mixed_poisson(custom_transfer=False):
 
     m = UnitSquareMesh(10, 10)
     nlevels = 2
@@ -14,7 +14,7 @@ def get_p1_space():
         return FunctionSpace(mesh, "CG", 1)
 
     def get_p1_prb_bcs():
-        return DirichletBC(get_p1_space(), Constant(0.0), "on_boundary")
+        return DirichletBC(get_p1_space(), 0, "on_boundary")
 
     def p1_callback():
         P1 = get_p1_space()
@@ -57,13 +57,28 @@ def p1_callback():
     appctx = {'get_coarse_operator': p1_callback,
               'get_coarse_space': get_p1_space,
               'coarse_space_bcs': get_p1_prb_bcs()}
-
-    solve(a == L, w, solver_parameters=params, appctx=appctx)
+    if custom_transfer:
+        P1 = get_p1_space()
+        V = FunctionSpace(mesh, "DGT", degree - 1)
+        I = assemble(Interpolate(TrialFunction(P1), V)).petscmat
+        R = PETSc.Mat().createTranspose(I)
+        appctx['interpolation_matrix'] = I
+        appctx['restriction_matrix'] = R
+
+    problem = LinearVariationalProblem(a, L, w)
+    solver = LinearVariationalSolver(problem, solver_parameters=params, appctx=appctx)
+    solver.solve()
     _, uh = w.subfunctions
 
     # Analytical solution
     f.interpolate(x[0]*(1-x[0])*x[1]*(1-x[1]))
 
+    if custom_transfer:
+        hyb = solver.snes.ksp.pc.getPythonContext()
+        gtmg = hyb.trace_ksp.pc.getPythonContext()
+        assert I.handle != R.handle
+        assert gtmg.pc.getMGInterpolation(1).handle == I.handle
+        assert gtmg.pc.getMGRestriction(1).handle == R.handle
     return errornorm(f, uh, norm_type="L2")
 
 
@@ -144,8 +159,9 @@ def p1_callback():
 
 
 @pytest.mark.skipcomplexnoslate
-def test_mixed_poisson_gtmg():
-    assert run_gtmg_mixed_poisson() < 1e-5
+@pytest.mark.parametrize("custom_transfer", [False, True])
+def test_mixed_poisson_gtmg(custom_transfer):
+    assert run_gtmg_mixed_poisson(custom_transfer) < 1e-5
 
 
 @pytest.mark.skipcomplexnoslate