Eikehmueller/gtmgpc restriction

firedrake/preconditioners/gtmg.py

@@ -1,3 +1,6 @@
+"""Non-nested multigrid preconditioner"""
+
+from firedrake_citations import Citations
 from firedrake.petsc import PETSc
 from firedrake.preconditioners.base import PCBase
 from firedrake.parameters import parameters
@@ -11,13 +14,60 @@
 
 
 class GTMGPC(PCBase):
+    """Non-nested multigrid preconditioner
+
+    Implements the method described in [1]. Note that while the authors of
+    this paper consider a non-nested function space hierarchy, the algorithm
+    can also be applied if the spaces are nested.
+
+    Uses PCMG to implement a two-level multigrid method involving two spaces:
+
+        * `V`: the fine space on which the problem is formulated
+        * `V_coarse`: a user-defined coarse space
+
+    The following options must be passed through the `appctx` dictionary:
+
+        * `get_coarse_space`: method which returns the user-defined coarse space
+        * `get_coarse_operator`: method which returns the operator on the coarse space
+
+    The following options (also passed through the `appctx`) are optional:
+
+        * `form_compiler_parameters`: parameters for assembling operators on
+          both levels of the hierarchy.
+        * `coarse_space_bcs`: boundary conditions to be used on coarse space.
+        * `get_coarse_op_nullspace`: method which returns the nullspace of the
+          coarse operator.
+        * `get_coarse_op_transpose_nullspace`: method which returns the
+          nullspace of the transpose of the coarse operator.
+        * `interpolation_matrix`: PETSc Mat which describes the interpolation
+          from the coarse to the fine space. If omitted, this will be
+          constructed automatically with an :class:`.Interpolate` object.
+        * `restriction_matrix`: PETSc Mat which describes the restriction
+          from the fine space dual to the coarse space dual. It defaults
+          to the transpose of the interpolation matrix.
+
+    PETSc options for the underlying PCMG object can be set with the
+    prefix ``gt_``.
+
+    Reference:
+
+        [1] Gopalakrishnan, J. and Tan, S., 2009: "A convergent multigrid
+        cycle for the hybridized mixed method". Numerical Linear Algebra
+        with Applications, 16(9), pp.689-714. https://doi.org/10.1002/nla.636
+
+    """
 
     needs_python_pmat = False
     _prefix = "gt_"
 
     def initialize(self, pc):
+        """Initialize new instance
+
+        :arg pc: PETSc preconditioner instance
+        """
         from firedrake.assemble import assemble, get_assembler
 
+        Citations().register("Gopalakrishnan2009")
         _, P = pc.getOperators()
         appctx = self.get_appctx(pc)
         fcp = appctx.get("form_compiler_parameters")
@@ -58,7 +108,7 @@ def initialize(self, pc):
         else:
             fine_petscmat = P
 
-        # Transfer fine operator null space
+        # Transfer fine operator nullspace
         fine_petscmat.setNullSpace(P.getNullSpace())
         fine_transpose_nullspace = P.getTransposeNullSpace()
         if fine_transpose_nullspace.handle != 0:
@@ -107,6 +157,7 @@ def initialize(self, pc):
             coarse_test, coarse_trial = coarse_operator.arguments()
             interp = assemble(Interpolate(coarse_trial, fine_space))
             interp_petscmat = interp.petscmat
+        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.
@@ -119,6 +170,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()
@@ -147,23 +200,46 @@ def initialize(self, pc):
             coarse_solver.setFromOptions()
 
     def update(self, pc):
+        """Update preconditioner
+
+        Re-assemble operators on both levels of the hierarchy
+
+        :arg pc: PETSc preconditioner instance
+        """
         if hasattr(self, "fine_op"):
             self._assemble_fine_op(tensor=self.fine_op)
 
         self._assemble_coarse_op(tensor=self.coarse_op)
         self.pc.setUp()
 
     def apply(self, pc, X, Y):
+        """Apply preconditioner
+
+        :arg pc: PETSc preconditioner instance
+        :arg X: right hand side PETSc vector
+        :arg Y: PETSc vector with resulting solution
+        """
         dm = self._dm
         with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
             self.pc.apply(X, Y)
 
     def applyTranspose(self, pc, X, Y):
+        """Apply transpose preconditioner
+
+        :arg pc: PETSc preconditioner instance
+        :arg X: right hand side PETSc vector
+        :arg Y: PETSc vector with resulting solution
+        """
         dm = self._dm
         with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
             self.pc.applyTranspose(X, Y)
 
     def view(self, pc, viewer=None):
+        """View preconditioner options
+
+        :arg pc: preconditioner instance
+        :arg viewer: PETSc viewer instance
+        """
         super(GTMGPC, self).view(pc, viewer)
         if hasattr(self, "pc"):
             viewer.printfASCII("PC using Gopalakrishnan and Tan algorithm\n")

firedrake_citations/__init__.py

@@ -292,6 +292,22 @@ def print_at_exit(cls):
 }
 """)
 
+Citations().add("Gopalakrishnan2009", """
+@article{Gopalakrishnan2009,
+  author    = {Jayadeep Gopalakrishnan and Shuguang Tan},
+  doi       = {10.1002/nla.636},
+  journal   = {Numerical Linear Algebra with Applications},
+  month     = {sep},
+  number    = {9},
+  pages     = {689--714},
+  publisher = {Wiley},
+  title     = {A convergent multigrid cycle for the hybridized mixed method},
+  url       = {https://doi.org/10.1002/nla.636},
+  volume    = {16},
+  year      = {2009}
+}
+""")
+
 Citations().add("nixonhill2023consistent", """
 @article{nixonhill2023consistent,
   author        = {Nixon-Hill, R. W. and Shapero, D. and Cotter, C. J. and Ham, D. A.},

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