Add in riesz map and remove pc update (#713)

Co-authored-by: Josh Hope-Collins <joshua.hope-collins13@imperial.ac.uk>

Author: 2 people

gusto/solvers/preconditioners.py

@@ -2,7 +2,7 @@
 
 from firedrake import (dot, jump, dS_h, ds_b, ds_t, ds,
                        FacetNormal, Tensor, AssembledVector,
-                       AuxiliaryOperatorPC, PETSc)
+                       AuxiliaryOperatorPC, PETSc, RieszMap)
 
 from firedrake.preconditioners import PCBase
 from firedrake.matrix_free.operators import ImplicitMatrixContext
@@ -460,35 +460,6 @@ class CompressibleHybridisedSCPC(PCBase):
 
     _prefix = "compressible_hybrid_scpc"
 
-    scpc_parameters = {
-        'mat_type': 'matfree',
-        'ksp_type': 'preonly',
-        'pc_type': 'python',
-        'pc_python_type': 'firedrake.SCPC',
-        'pc_sc_eliminate_fields': '0, 1',
-        # The reduced operator is not symmetric
-        'condensed_field': {
-            'ksp_type': 'fgmres',
-            'ksp_rtol': 1.0e-8,
-            'ksp_atol': 1.0e-8,
-            'ksp_max_it': 100,
-            'pc_type': 'gamg',
-            'pc_gamg_sym_graph': None,
-            'mg_levels': {
-                'ksp_type': 'gmres',
-                'ksp_max_it': 5,
-                'pc_type': 'bjacobi',
-                'sub_pc_type': 'ilu'
-            }
-        }
-    }
-
-    cg_ilu_parameters = {
-        'ksp_type': 'cg',
-        'pc_type': 'bjacobi',
-        'sub_pc_type': 'ilu'
-    }
-
     def initialize(self, pc):
         """
         Set up problem and solver
@@ -507,15 +478,14 @@ def initialize(self, pc):
             raise ValueError("Expecting PC type python")
 
         self._process_context(pc)
+        prefix = pc.getOptionsPrefix()
 
-        self.hybridised_scpc_parameters = self.scpc_parameters
-        self.rho_avg_solver_parameters = self.cg_ilu_parameters
-        self.theta_solver_parameters = self.cg_ilu_parameters
-        self.exner_avg_solver_parameters = self.cg_ilu_parameters
+        # Unpack sub-solver parameters from the options prefix
+        self.riesz_map_parameters = PETSc.Options(prefix + "riesz_map_").getAll()
 
         if logger.isEnabledFor(DEBUG):
-            self.hybridised_scpc_parameters['ksp_monitor_true_residual'] = None
-            self.hybridised_scpc_parameters['ksp_converged_reason'] = None
+            self.scpc_solve_parameters['ksp_monitor_true_residual'] = None
+            self.scpc_solve_parameters['ksp_converged_reason'] = None
 
         # Equations and parameters
         equations = self.equations
@@ -548,6 +518,11 @@ def initialize(self, pc):
         self.xrhs = Function(self.W)
         self.y = Function(self.W)
 
+        # Riesz map for the dual of the original function space
+        self.riesz_map = RieszMap(self.W.dual(),
+                                  solver_parameters=self.riesz_map_parameters,
+                                  constant_jacobian=True)
+
         self.y_hybrid = Function(self.W_hyb)
 
         # Split input functions
@@ -632,12 +607,12 @@ def L_tr(f):
                                                  constant_jacobian=True)
 
         self.rho_avg_solver = LinearVariationalSolver(
-            rho_avg_prb, solver_parameters=self.rho_avg_solver_parameters,
-            options_prefix=pc.getOptionsPrefix()+'rhobar_avg_solver'
+            rho_avg_prb,
+            options_prefix=pc.getOptionsPrefix()+'rhobar_avg'
         )
         self.exner_avg_solver = LinearVariationalSolver(
-            exner_avg_prb, solver_parameters=self.exner_avg_solver_parameters,
-            options_prefix=pc.getOptionsPrefix()+'exnerbar_avg_solver'
+            exner_avg_prb,
+            options_prefix=pc.getOptionsPrefix()+'exnerbar_avg'
         )
 
         # "broken" u, rho, and trace system
@@ -691,8 +666,9 @@ def L_tr(f):
             aeqn, Leqn, self.y_hybrid, constant_jacobian=True
         )
         self.hybridized_solver = LinearVariationalSolver(
-            hybridized_prb, solver_parameters=self.hybridised_scpc_parameters,
-            options_prefix=pc.getOptionsPrefix()+self._prefix, appctx=appctx
+            hybridized_prb,
+            options_prefix=pc.getOptionsPrefix()+self._prefix,
+            appctx=appctx
         )
 
         if logger.isEnabledFor(DEBUG):
@@ -725,8 +701,8 @@ def L_tr(f):
         )
 
         self.theta_solver = LinearVariationalSolver(
-            theta_problem, solver_parameters=self.theta_solver_parameters,
-            options_prefix=pc.getOptionsPrefix()+'thetabacksubstitution'
+            theta_problem,
+            options_prefix=pc.getOptionsPrefix()+'theta_backsub'
         )
 
         if logger.isEnabledFor(DEBUG):
@@ -753,7 +729,7 @@ def apply(self, pc, x, y):
         with self.xstar.dat.vec_wo as xv:
             x.copy(xv)
 
-        self.xrhs.assign(self.xstar.riesz_representation())
+        self.xrhs.assign(self.riesz_map(self.xstar))
         # Solve hybridized system
         self.hybridized_solver.solve()
 

gusto/solvers/solver_presets.py

@@ -67,12 +67,69 @@ def trace_nullsp(T):
         # Compressible Euler equations - (u, rho, theta) system. We use a
         # bespoke hybridization preconditioner for this system owing to the
         # nonlinear pressure gradient term.
+        r_tol = 1e-8
+
+        theta_backsub_settings = {
+            'ksp_type': 'cg',
+            'pc_type': 'bjacobi',
+            'sub_pc_type': 'ilu',
+            'ksp_rtol': r_tol
+        }
+
+        exnerbar_avg_settings = {
+            'ksp_type': 'cg',
+            'pc_type': 'bjacobi',
+            'sub_pc_type': 'ilu',
+            'ksp_rtol': r_tol
+        }
+
+        rhobar_avg_settings = {
+            'ksp_type': 'cg',
+            'pc_type': 'bjacobi',
+            'sub_pc_type': 'ilu',
+            'ksp_rtol': r_tol
+        }
+
+        riesz_map_settings = {
+            'ksp_type': 'cg',
+            'pc_type': 'bjacobi',
+            'sub_pc_type': 'ilu',
+            'ksp_rtol': r_tol
+        }
+
+        scpc_solve_settings = {
+            'mat_type': 'matfree',
+            'ksp_type': 'preonly',
+            'pc_type': 'python',
+            'pc_python_type': 'firedrake.SCPC',
+            'pc_sc_eliminate_fields': '0,1',
+            # The reduced operator is not symmetric
+            'condensed_field': {
+                'ksp_type': 'fgmres',
+                'ksp_rtol': r_tol,
+                'ksp_max_it': 100,
+                'pc_type': 'gamg',
+                'pc_gamg_sym_graph': None,
+                'mg_levels': {
+                    'ksp_type': 'gmres',
+                    'ksp_max_it': 5,
+                    'pc_type': 'bjacobi',
+                    'sub_pc_type': 'ilu'
+                }
+            }
+        }
+
         settings = {
             'ksp_monitor': None,
             'ksp_type': 'preonly',
             'mat_type': 'matfree',
             'pc_type': 'python',
             'pc_python_type': 'gusto.CompressibleHybridisedSCPC',
+            'theta_backsub': theta_backsub_settings,
+            'exnerbar_avg': exnerbar_avg_settings,
+            'rhobar_avg': rhobar_avg_settings,
+            'riesz_map': riesz_map_settings,
+            'compressible_hybrid_scpc': scpc_solve_settings,
         }
 
         # We pass the implicit weighting parameter (alpha) and tau_values to the

gusto/timestepping/semi_implicit_quasi_newton.py

@@ -20,7 +20,6 @@
 from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation
 from gusto.timestepping.timestepper import BaseTimestepper
 from gusto.solvers.solver_presets import hybridised_solver_parameters
-from gusto.equations.compressible_euler_equations import CompressibleEulerEquations
 
 
 __all__ = ["SemiImplicitQuasiNewton", "Forcing", "QuasiNewtonLinearSolver"]
@@ -937,12 +936,6 @@ def update_reference_profiles(self):
             self.equation.update_reference_profiles()
             self.solver.invalidate_jacobian()
 
-            # TODO: Issue #686 is to address this reference profile update bug (pythonPC update not called)
-            # this line forces it to update for now
-            pc = self.solver.snes.getKSP().getPC()
-            if (isinstance(self.equation, CompressibleEulerEquations) and pc.getType() == "python"):
-                pc.getPythonContext().update(pc)
-
     def zero_non_prognostics(self, equation, xrhs, field_names, prognostic_names):
         """
         Zero rhs term F(x) for non-prognostics.