Dwitt/tr bdf2 (#643)

Co-authored-by: jshipton <j.shipton@exeter.ac.uk>

Author: Witt-D

examples/compressible_euler/skamarock_klemp_nonhydrostatic.py

@@ -18,15 +18,15 @@
 import itertools
 from firedrake import (
     as_vector, SpatialCoordinate, PeriodicIntervalMesh, ExtrudedMesh, exp, sin,
-    Function, pi, COMM_WORLD
+    Function, pi, COMM_WORLD, sqrt
 )
 import numpy as np
 from gusto import (
-    Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
-    logger, SUPGOptions, Perturbation, CompressibleParameters,
+    Domain, IO, OutputParameters, TRBDF2QuasiNewton, SemiImplicitQuasiNewton, SSPRK3,
+    DGUpwind, logger, SUPGOptions, Perturbation, CompressibleParameters,
     CompressibleEulerEquations, HydrostaticCompressibleEulerEquations,
     compressible_hydrostatic_balance, RungeKuttaFormulation, CompressibleSolver,
-    SubcyclingOptions, hydrostatic_parameters
+    hydrostatic_parameters, SubcyclingOptions,
 )
 
 skamarock_klemp_nonhydrostatic_defaults = {
@@ -36,7 +36,8 @@
     'tmax': 3000.,
     'dumpfreq': 250,
     'dirname': 'skamarock_klemp_nonhydrostatic',
-    'hydrostatic': False
+    'hydrostatic': False,
+    'timestepper': 'SIQN'
 }
 
 
@@ -47,7 +48,8 @@ def skamarock_klemp_nonhydrostatic(
         tmax=skamarock_klemp_nonhydrostatic_defaults['tmax'],
         dumpfreq=skamarock_klemp_nonhydrostatic_defaults['dumpfreq'],
         dirname=skamarock_klemp_nonhydrostatic_defaults['dirname'],
-        hydrostatic=skamarock_klemp_nonhydrostatic_defaults['hydrostatic']
+        hydrostatic=skamarock_klemp_nonhydrostatic_defaults['hydrostatic'],
+        timestepper=skamarock_klemp_nonhydrostatic_defaults['timestepper']
 ):
 
     # ------------------------------------------------------------------------ #
@@ -98,7 +100,7 @@ def skamarock_klemp_nonhydrostatic(
         output = OutputParameters(
             dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq,
             dump_vtus=False, dump_nc=True,
-            point_data=[('theta_perturbation', points)],
+            point_data=[('theta_perturbation', points)]
         )
     else:
         logger.warning(
@@ -115,7 +117,10 @@ def skamarock_klemp_nonhydrostatic(
 
     # Transport schemes
     theta_opts = SUPGOptions()
-    subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25)
+    if timestepper == 'SIQN':
+        subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25)
+    else:
+        subcycling_options = None
     transported_fields = [
         SSPRK3(domain, "u", subcycling_options=subcycling_options),
         SSPRK3(
@@ -135,6 +140,9 @@ def skamarock_klemp_nonhydrostatic(
 
     # Linear solver
     if hydrostatic:
+        if timestepper == 'TR-BDF2':
+            raise ValueError('Hydrostatic equations not implmented for TR-BDF2')
+
         linear_solver = CompressibleSolver(
             eqns, solver_parameters=hydrostatic_parameters,
             overwrite_solver_parameters=True
@@ -143,11 +151,23 @@ def skamarock_klemp_nonhydrostatic(
         linear_solver = CompressibleSolver(eqns)
 
     # Time stepper
-    stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver
-    )
+    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,
+            tr_solver=tr_solver,
+            bdf_solver=bdf_solver
+        )
 
+    elif timestepper == 'SIQN':
+        stepper = SemiImplicitQuasiNewton(
+            eqns, io, transported_fields, transport_methods,
+            linear_solver=linear_solver
+        )
     # ------------------------------------------------------------------------ #
     # Initial conditions
     # ------------------------------------------------------------------------ #
@@ -248,6 +268,13 @@ def skamarock_klemp_nonhydrostatic(
         action="store_true",
         default=skamarock_klemp_nonhydrostatic_defaults['hydrostatic']
     )
+    parser.add_argument(
+        'timestepper',
+        help='Which time stepper to use, takes SIQN or TR-BDF2',
+        type=str,
+        choices=['SIQN', 'TR-BDF2'],
+        default=skamarock_klemp_nonhydrostatic_defaults['timestepper']
+    )
     args, unknown = parser.parse_known_args()
 
     skamarock_klemp_nonhydrostatic(**vars(args))

examples/compressible_euler/test_compressible_euler_examples.py

@@ -83,6 +83,26 @@ def test_skamarock_klemp_nonhydrostatic_parallel():
     test_skamarock_klemp_nonhydrostatic()
 
 
+def test_skamarock_klemp_nonhydrostatic_TR_BDF2():
+    from skamarock_klemp_nonhydrostatic import skamarock_klemp_nonhydrostatic
+    test_name = 'skamarock_klemp_nonhydrostatic'
+    skamarock_klemp_nonhydrostatic(
+        ncolumns=30,
+        nlayers=5,
+        dt=6.0,
+        tmax=60.0,
+        dumpfreq=10,
+        dirname=make_dirname(test_name),
+        hydrostatic=False,
+        timestepper='TR-BDF2'
+    )
+
+
+@pytest.mark.parallel(nprocs=2)
+def test_skamarock_klemp_nonhydrostatic_TR_BDF2_parallel():
+    test_skamarock_klemp_nonhydrostatic_TR_BDF2()
+
+
 # Hydrostatic equations not currently working
 @pytest.mark.xfail
 def test_hyd_switch_skamarock_klemp_nonhydrostatic():

gusto/timestepping/__init__.py

@@ -1,3 +1,4 @@
 from gusto.timestepping.timestepper import *                 # noqa
 from gusto.timestepping.split_timestepper import *           # noqa
-from gusto.timestepping.semi_implicit_quasi_newton import *  # noqa
+from gusto.timestepping.semi_implicit_quasi_newton import *  # noqa
+from gusto.timestepping.tr_bdf2_quasi_newton import *        # noqa

gusto/timestepping/semi_implicit_quasi_newton.py

@@ -22,7 +22,7 @@
 from gusto.timestepping.timestepper import BaseTimestepper
 
 
-__all__ = ["SemiImplicitQuasiNewton"]
+__all__ = ["SemiImplicitQuasiNewton", "Forcing"]
 
 
 class SemiImplicitQuasiNewton(BaseTimestepper):
@@ -532,7 +532,6 @@ def timestep(self):
 
             xrhs.assign(0.)  # xrhs is the residual which goes in the linear solve
             xrhs_phys.assign(x_after_fast(self.field_name) - xp(self.field_name))
-
             for inner in range(self.num_inner):
 
                 # Implicit forcing ---------------------------------------------
@@ -556,7 +555,6 @@ def timestep(self):
 
             # Update xnp1 values for active tracers not included in the linear solve
             self.copy_active_tracers(x_after_fast, xnp1)
-
             self._apply_bcs()
 
         for name, scheme in self.auxiliary_schemes:
@@ -611,7 +609,7 @@ class Forcing(object):
     semi-implicit time discretisation.
     """
 
-    def __init__(self, equation, implicit_terms, alpha):
+    def __init__(self, equation, implicit_terms, alpha, dt=None):
         """
         Args:
             equation (:class:`PrognosticEquationSet`): the prognostic equations
@@ -621,10 +619,14 @@ def __init__(self, equation, implicit_terms, alpha):
             alpha (:class:`Function`): semi-implicit off-centering factor. An
                 alpha of 0 corresponds to fully explicit, while a factor of 1
                 corresponds to fully implicit.
+            dt (:float): timestep over which to apply forcing, defaults to None
+                in which case it is taken from the equation class.
         """
 
         self.field_name = equation.field_name
-        dt = equation.domain.dt
+
+        if dt is None:
+            dt = equation.domain.dt
 
         W = equation.function_space
         self.x0 = Function(W)
@@ -645,16 +647,14 @@ def __init__(self, equation, implicit_terms, alpha):
                                replace_subject(trials),
                                map_if_false=drop)
 
-        # the explicit forms are multiplied by (1-alpha) and moved to the rhs
-        one_minus_alpha = Function(alpha.function_space(), val=1-alpha)
-        L_explicit = -one_minus_alpha*dt*residual.label_map(
+        L_explicit = -(1 - alpha)*dt*residual.label_map(
             lambda t:
                 any(t.has_label(time_derivative, hydrostatic, *implicit_terms,
                                 return_tuple=True)),
             drop,
             replace_subject(self.x0))
 
-        # the implicit forms are multiplied by alpha and moved to the rhs
+        # the implicit forms are multiplied by implicit scaling and moved to the rhs
         L_implicit = -alpha*dt*residual.label_map(
             lambda t:
                 any(t.has_label(
@@ -682,10 +682,11 @@ def __init__(self, equation, implicit_terms, alpha):
                 drop)
 
         # now we can set up the explicit and implicit problems
-        explicit_forcing_problem = LinearVariationalProblem(
-            a.form, L_explicit.form, self.xF, bcs=bcs,
-            constant_jacobian=True
-        )
+        if alpha != 1.0:
+            explicit_forcing_problem = LinearVariationalProblem(
+                a.form, L_explicit.form, self.xF, bcs=bcs,
+                constant_jacobian=True
+            )
 
         implicit_forcing_problem = LinearVariationalProblem(
             a.form, L_implicit.form, self.xF, bcs=bcs,
@@ -695,19 +696,21 @@ def __init__(self, equation, implicit_terms, alpha):
         self.solver_parameters = mass_parameters(W, equation.domain.spaces)
 
         self.solvers = {}
-        self.solvers["explicit"] = LinearVariationalSolver(
-            explicit_forcing_problem,
-            solver_parameters=self.solver_parameters,
-            options_prefix="ExplicitForcingSolver"
-        )
+        if alpha != 1.0:
+            self.solvers["explicit"] = LinearVariationalSolver(
+                explicit_forcing_problem,
+                solver_parameters=self.solver_parameters,
+                options_prefix="ExplicitForcingSolver"
+            )
         self.solvers["implicit"] = LinearVariationalSolver(
             implicit_forcing_problem,
             solver_parameters=self.solver_parameters,
             options_prefix="ImplicitForcingSolver"
         )
 
         if logger.isEnabledFor(DEBUG):
-            self.solvers["explicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual)
+            if alpha != 1.0:
+                self.solvers["explicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual)
             self.solvers["implicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual)
 
     def apply(self, x_in, x_nl, x_out, label):