Generalise linear solvers

examples/boussinesq/skamarock_klemp_compressible.py

@@ -15,7 +15,7 @@
 from gusto import (
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, SUPGOptions, Divergence, Perturbation, CourantNumber,
-    BoussinesqParameters, BoussinesqEquations, BoussinesqSolver,
+    BoussinesqParameters, BoussinesqEquations, BoussinesqSolver, LinearTimesteppingSolver,
     boussinesq_hydrostatic_balance
 )
 
@@ -91,7 +91,21 @@ def skamarock_klemp_compressible_bouss(
     ]
 
     # Linear solver
-    linear_solver = BoussinesqSolver(eqns)
+    solver_parameters = HybridisedSolverParameters(eqns.name)
+    def trace_nullsp(T):
+        return VectorSpaceBasis(constant=True)
+
+
+    appctx = {
+        'auxform': eqns.schur_complement_form(alpha=0.5),
+        "trace_nullspace": trace_nullsp,
+    }
+
+    linear_solver = SIQNLinearSolver(
+        eqns, solver_prognostics=["u", "p", "b"], alpha=0.5, implicit_terms=[incompressible, sponge],
+        solver_parameters=solver_parameters,
+        appctx=appctx
+    )
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(

examples/boussinesq/skamarock_klemp_incompressible.py

@@ -16,7 +16,8 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, SUPGOptions, Divergence, Perturbation, CourantNumber,
     BoussinesqParameters, BoussinesqEquations, BoussinesqSolver,
-    boussinesq_hydrostatic_balance
+    boussinesq_hydrostatic_balance, incompressible, sponge, SIQNLinearSolver,
+    HybridisedSolverParameters, LinearTimesteppingSolver
 )
 
 skamarock_klemp_incompressible_bouss_defaults = {
@@ -88,7 +89,28 @@ def skamarock_klemp_incompressible_bouss(
     ]
 
     # Linear solver
-    linear_solver = BoussinesqSolver(eqns)
+    solver_parameters = HybridisedSolverParameters(eqns.name)
+    def trace_nullsp(T):
+        return VectorSpaceBasis(constant=True)
+
+
+    appctx = {
+        'auxform': eqns.schur_complement_form(alpha=0.5),
+        "trace_nullspace": trace_nullsp,
+    }
+
+    linear_solver = SIQNLinearSolver(
+        eqns, solver_prognostics=["u", "p", "b"], alpha=0.5, implicit_terms=[incompressible, sponge],
+        solver_parameters=solver_parameters,
+        appctx=appctx
+    )
+    linear_solver = LinearTimesteppingSolver(eqns, alpha=0.5, options_prefix="boussinesq",
+        solver_parameters=solver_parameters,
+        appctx=appctx
+    )
+
+    #linear_solver = BoussinesqSolver(eqns)
+
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(

examples/compressible_euler/dcmip_3_1_gravity_wave.py

@@ -15,9 +15,9 @@
 from gusto import (
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, SUPGOptions, lonlatr_from_xyz, CompressibleParameters,
-    CompressibleEulerEquations, CompressibleSolver, ZonalComponent,
+    CompressibleEulerEquations, ZonalComponent,
     compressible_hydrostatic_balance, Perturbation, GeneralCubedSphereMesh,
-    SubcyclingOptions
+    SubcyclingOptions, HybridisedSolverParameters, SIQNLinearSolver, incompressible, sponge
 )
 
 dcmip_3_1_gravity_wave_defaults = {
@@ -105,13 +105,28 @@ def dcmip_3_1_gravity_wave(
         DGUpwind(eqns, field) for field in ["u", "rho", "theta"]
     ]
 
-    # Linear solver
-    linear_solver = CompressibleSolver(eqns)
+    # # Linear solver
+    # solver_parameters = HybridisedSolverParameters(eqns.name)
+
+    # def trace_nullsp(T):
+    #     return VectorSpaceBasis(constant=True)
+
+
+    # appctx = {
+    #     'equations': eqns,
+    #     'alpha': 0.5,
+    #     'trace_nullspace': trace_nullsp
+    # }
+
+    # linear_solver = SIQNLinearSolver(
+    #     eqns, solver_prognostics=["u", "rho", "theta"], alpha=0.5, implicit_terms=[incompressible, sponge],
+    #     solver_parameters=solver_parameters,
+    #     appctx=appctx, enforce_pc_on_rhs=True
+    # )
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver
+        eqns, io, transported_fields, transport_methods
     )
 
     # ------------------------------------------------------------------------ #

examples/compressible_euler/dry_bryan_fritsch.py

@@ -15,8 +15,9 @@
 from gusto import (
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     RecoverySpaces, BoundaryMethod, Perturbation, CompressibleParameters,
-    CompressibleEulerEquations, CompressibleSolver,
-    compressible_hydrostatic_balance
+    CompressibleEulerEquations,
+    compressible_hydrostatic_balance, HybridisedSolverParameters, SIQNLinearSolver,
+    incompressible, sponge
 )
 
 dry_bryan_fritsch_defaults = {
@@ -102,13 +103,9 @@ def dry_bryan_fritsch(
         DGUpwind(eqns, field) for field in ["u", "rho", "theta"]
     ]
 
-    # Linear solver
-    linear_solver = CompressibleSolver(eqns)
-
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver
+        eqns, io, transported_fields, transport_methods
     )
 
     # ------------------------------------------------------------------------ #

examples/compressible_euler/schaer_mountain.py

@@ -19,7 +19,8 @@
     OutputParameters, IO, SSPRK3, DGUpwind, SemiImplicitQuasiNewton,
     compressible_hydrostatic_balance, SpongeLayerParameters, Exner, ZComponent,
     Perturbation, SUPGOptions, TrapeziumRule, MaxKernel, MinKernel,
-    CompressibleEulerEquations, SubcyclingOptions, RungeKuttaFormulation
+    CompressibleEulerEquations, SubcyclingOptions, RungeKuttaFormulation,
+    HybridisedSolverParameters, SIQNLinearSolver, incompressible, sponge
 )
 
 schaer_mountain_defaults = {
@@ -97,11 +98,11 @@ def schaer_mountain(
 
     # Equation
     parameters = CompressibleParameters(mesh, g=g, cp=cp)
-    sponge = SpongeLayerParameters(
+    sponge_params = SpongeLayerParameters(
         mesh, H=domain_height, z_level=domain_height-sponge_depth, mubar=mu_dt/dt
     )
     eqns = CompressibleEulerEquations(
-        domain, parameters, sponge_options=sponge, u_transport_option=u_eqn_type
+        domain, parameters, sponge_options=sponge_params, u_transport_option=u_eqn_type
     )
 
     # I/O
@@ -134,14 +135,10 @@ def schaer_mountain(
         DGUpwind(eqns, "theta", ibp=theta_opts.ibp)
     ]
 
-    # Linear solver
-    tau_values = {'rho': 1.0, 'theta': 1.0}
-    linear_solver = CompressibleSolver(eqns, alpha, tau_values=tau_values)
-
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
         eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver, alpha=alpha, spinup_steps=spinup_steps
+        alpha=alpha, spinup_steps=spinup_steps
     )
 
     # ------------------------------------------------------------------------ #

examples/compressible_euler/skamarock_klemp_nonhydrostatic.py

@@ -13,21 +13,21 @@
 """
 from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
 
-from petsc4py import PETSc
-PETSc.Sys.popErrorHandler()
 import itertools
 from firedrake import (
     as_vector, SpatialCoordinate, PeriodicIntervalMesh, ExtrudedMesh, exp, sin,
-    Function, pi, COMM_WORLD, sqrt
+    PETSc, Function, pi, COMM_WORLD, sqrt
 )
 import numpy as np
 from gusto import (
     Domain, IO, OutputParameters, TRBDF2QuasiNewton, SemiImplicitQuasiNewton, SSPRK3,
     DGUpwind, logger, SUPGOptions, Perturbation, CompressibleParameters,
     CompressibleEulerEquations, HydrostaticCompressibleEulerEquations,
-    compressible_hydrostatic_balance, RungeKuttaFormulation, CompressibleSolver,
-    hydrostatic_parameters, SubcyclingOptions,
+    compressible_hydrostatic_balance, RungeKuttaFormulation,
+    SubcyclingOptions, incompressible, sponge, SIQNLinearSolver,
+    HybridisedSolverParameters
 )
+PETSc.Sys.popErrorHandler()
 
 skamarock_klemp_nonhydrostatic_defaults = {
     'ncolumns': 150,
@@ -138,24 +138,29 @@ def skamarock_klemp_nonhydrostatic(
         DGUpwind(eqns, "theta", ibp=theta_opts.ibp)
     ]
 
-    # Linear solver
-    if hydrostatic:
-        if timestepper == 'TR-BDF2':
+    # Time stepper and linear solver
+    if timestepper == 'TR-BDF2':
+        if hydrostatic:
             raise ValueError('Hydrostatic equations not implmented for TR-BDF2')
+        gamma = (1-sqrt(2)/2)
+        gamma2 = (1 - 2*float(gamma))/(2 - 2*float(gamma))
+
 
-        linear_solver = CompressibleSolver(
-            eqns, solver_parameters=hydrostatic_parameters,
-            overwrite_solver_parameters=True
+        # Linear solver
+        solver_parameters_tr, appctx_tr = HybridisedSolverParameters(eqns, alpha=gamma)
+        solver_parameters_bdf, appctx_bdf = HybridisedSolverParameters(eqns, alpha=gamma2)
+
+        tr_solver = SIQNLinearSolver(
+            eqns, solver_prognostics=["u", "rho", "theta"], alpha=gamma, implicit_terms=[incompressible, sponge],
+            solver_parameters=solver_parameters_tr,
+            appctx=appctx_tr, enforce_pc_on_rhs=True
+        )
+        bdf_solver = SIQNLinearSolver(
+            eqns, solver_prognostics=["u", "rho", "theta"], alpha=gamma2, implicit_terms=[incompressible, sponge],
+            solver_parameters=solver_parameters_bdf,
+            appctx=appctx_bdf, enforce_pc_on_rhs=True
         )
-    else:
-        linear_solver = CompressibleSolver(eqns)
 
-    # Time stepper
-    if timestepper == 'TR-BDF2':
-        gamma = (1-sqrt(2)/2)
-        gamma2 = (1 - 2*float(gamma))/(2 - 2*float(gamma))
-        tr_solver = CompressibleSolver(eqns, alpha=gamma)
-        bdf_solver = CompressibleSolver(eqns, alpha=gamma2)
         stepper = TRBDF2QuasiNewton(
             eqns, io, transported_fields, transport_methods,
             gamma=gamma,
@@ -165,8 +170,7 @@ def skamarock_klemp_nonhydrostatic(
 
     elif timestepper == 'SIQN':
         stepper = SemiImplicitQuasiNewton(
-            eqns, io, transported_fields, transport_methods,
-            linear_solver=linear_solver
+            eqns, io, transported_fields, transport_methods
         )
     # ------------------------------------------------------------------------ #
     # Initial conditions