User-defined preconditioning of local solver in HybridisationPC

firedrake/slate/static_condensation/hybridization.py

@@ -213,10 +213,15 @@ def initialize(self, pc):
         K1 = Tensor(split_trace_op[(0, id1)])
         list_split_trace_ops = [K0, K1]
 
+        # Get the schur complement approximation
+        test, trial = A11.arguments()
+        schur_approx = self.get_user_schur_approx(pc, test, trial)
+
         # Build schur complement operator and right hand side
         nested = PETSc.Options(prefix).getBool("nested_schur", False)
         schur_rhs, schur_comp = self.build_schur(Atilde, K, list_split_mixed_ops,
-                                                 list_split_trace_ops, nested=nested)
+                                                 list_split_trace_ops, nested=nested,
+                                                 schur_approx=schur_approx)
 
         # Assemble the Schur complement operator and right-hand side
         self.schur_rhs = Function(TraceSpace)
@@ -285,9 +290,9 @@ def initialize(self, pc):
             trace_ksp.setFromOptions()
 
         # Generate reconstruction calls
-        self._reconstruction_calls(list_split_mixed_ops, list_split_trace_ops)
+        self._reconstruction_calls(list_split_mixed_ops, list_split_trace_ops, schur_approx)
 
-    def _reconstruction_calls(self, list_split_mixed_ops, list_split_trace_ops):
+    def _reconstruction_calls(self, list_split_mixed_ops, list_split_trace_ops, schur_approx=None):
         """This generates the reconstruction calls for the unknowns using the
         Lagrange multipliers.
 
@@ -316,10 +321,13 @@ def _reconstruction_calls(self, list_split_mixed_ops, list_split_trace_ops):
         u = split_sol[id1]
         lambdar = AssembledVector(self.trace_solution)
 
-        M = D - C * A.inv * B
+        schur = D - C * A.inv * B
+        schur = schur_approx.inv * schur if schur_approx else schur
         R = K_1.T - C * A.inv * K_0.T
-        u_rec = M.solve(f - C * A.inv * g - R * lambdar,
-                        decomposition="PartialPivLU")
+        rhs = f - C * A.inv * g - R * lambdar
+        rhs = schur_approx.inv * rhs if schur_approx else rhs
+
+        u_rec = schur.solve(rhs, decomposition="PartialPivLU")
         self._sub_unknown = functools.partial(assemble,
                                               u_rec,
                                               tensor=u,
@@ -334,7 +342,7 @@ def _reconstruction_calls(self, list_split_mixed_ops, list_split_trace_ops):
                                                form_compiler_parameters=self.ctx.fc_params,
                                                assembly_type="residual")
 
-    def build_schur(self, Atilde, K, list_split_mixed_ops, list_split_trace_ops, nested=False):
+    def build_schur(self, Atilde, K, list_split_mixed_ops, list_split_trace_ops, nested=False, schur_approx=None):
         """The Schur complement in the operators of the trace solve contains
         the inverse on a mixed system.  Users may want this inverse to be treated
         with another schur complement.
@@ -362,6 +370,8 @@ def build_schur(self, Atilde, K, list_split_mixed_ops, list_split_trace_ops, nes
 
             # inner schur complement
             S = (A11 - A10 * A00.inv * A01)
+            # preconditioning
+            S = schur_approx.inv * S if schur_approx else S
             # K * block1
             K_Ainv_block1 = [K0, -K0 * A00.inv * A01 + K1]
             # K * block1 * block2
@@ -377,6 +387,25 @@ def build_schur(self, Atilde, K, list_split_mixed_ops, list_split_trace_ops, nes
             schur_comp = K * Atilde.inv * K.T
         return schur_rhs, schur_comp
 
+    def get_user_schur_approx(self, pc, test, trial):
+        """Retrieve a user-defined AuxiliaryOperator from the PETSc Options,
+        which is an approximation to the Schur complement and its inverse is used
+        to precondition the local solve in the reconstruction calls (e.g.).
+        """
+        prefix_schur = pc.getOptionsPrefix() + "hybridization_approx_schur_"
+        sentinel = object()
+        usercode = PETSc.Options().getString(prefix_schur + 'pc_python_type', default=sentinel)
+
+        if usercode != sentinel:
+            (modname, funname) = usercode.rsplit('.', 1)
+            mod = __import__(modname)
+            fun = getattr(mod, funname)
+            if isinstance(fun, type):
+                fun = fun()
+            return Tensor(fun.form(pc, test, trial)[0])
+        else:
+            return None
+
     @PETSc.Log.EventDecorator("HybridUpdate")
     def update(self, pc):
         """Update by assembling into the operator. No need to

tests/slate/test_slate_hybridization.py

@@ -130,3 +130,65 @@ def test_slate_hybridization_nested_schur():
 
     assert sigma_err < 1e-11
     assert u_err < 1e-11
+
+
+class DGLaplacian(AuxiliaryOperatorPC):
+    def form(self, pc, u, v):
+        W = u.function_space()
+        n = FacetNormal(W.mesh())
+        alpha = Constant(3**3)
+        gamma = Constant(4**3)
+        h = CellSize(W.mesh())
+        h_avg = (h('+') + h('-'))/2
+        a_dg = -(inner(grad(u), grad(v))*dx
+                 - inner(jump(u, n), avg(grad(v)))*dS
+                 - inner(avg(grad(u)), jump(v, n), )*dS
+                 + alpha/h_avg * inner(jump(u, n), jump(v, n))*dS
+                 - inner(u*n, grad(v))*ds
+                 - inner(grad(u), v*n)*ds
+                 + (gamma/h)*inner(u, v)*ds)
+        bcs = None
+        return (a_dg, bcs)
+
+
+def test_mixed_poisson_approximated_schur():
+    # NOTE With the setup in this test, using the approximated schur complemement
+    # defined as DGLaplacian as a preconditioner to the schur complement in the
+    # reconstruction calls reduces the condition number of the local solve from
+    # 16.77 to 5.95
+    a, L, W = setup_poisson()
+
+    # Compare hybridized solution with non-hybridized
+    w = Function(W)
+    bcs = []
+
+    w = Function(W)
+    params = {'ksp_type': 'preonly',
+              'pc_type': 'python',
+              'mat_type': 'matfree',
+              'pc_python_type': 'firedrake.HybridizationPC',
+              'hybridization': {'ksp_type': 'preonly',
+                                'pc_type': 'lu',
+                                'approx_schur': {'ksp_type': 'preonly',
+                                                 'pc_type': 'python',
+                                                 'pc_python_type': __name__ + ".DGLaplacian"}}}
+
+    solve(a == L, w, bcs=bcs, solver_parameters=params)
+    sigma_h, u_h = w.split()
+
+    w2 = Function(W)
+    aij_params = {'ksp_type': 'preonly',
+                  'pc_type': 'python',
+                  'mat_type': 'matfree',
+                  'pc_python_type': 'firedrake.HybridizationPC',
+                  'hybridization': {'ksp_type': 'preonly',
+                                    'pc_type': 'lu'}}
+    solve(a == L, w2, bcs=bcs, solver_parameters=aij_params)
+    _sigma, _u = w2.split()
+
+    # Return the L2 error
+    sigma_err = errornorm(sigma_h, _sigma)
+    u_err = errornorm(u_h, _u)
+
+    assert sigma_err < 1e-8
+    assert u_err < 1e-8