Predictors (#551)

examples/shallow_water/williamson_5.py

@@ -52,6 +52,7 @@ def williamson_5(
     # ------------------------------------------------------------------------ #
 
     element_order = 1
+
     # ------------------------------------------------------------------------ #
     # Set up model objects
     # ------------------------------------------------------------------------ #

gusto/timestepping/semi_implicit_quasi_newton.py

@@ -3,8 +3,10 @@
 and GungHo dynamical cores.
 """
 
-from firedrake import (Function, Constant, TrialFunctions, DirichletBC,
-                       LinearVariationalProblem, LinearVariationalSolver)
+from firedrake import (
+    Function, Constant, TrialFunctions, DirichletBC, div, Interpolator,
+    LinearVariationalProblem, LinearVariationalSolver
+)
 from firedrake.fml import drop, replace_subject
 from pyop2.profiling import timed_stage
 from gusto.core import TimeLevelFields, StateFields
@@ -36,8 +38,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
                  slow_physics_schemes=None, fast_physics_schemes=None,
                  alpha=Constant(0.5), off_centred_u=False,
                  num_outer=2, num_inner=2, accelerator=False,
-                 reference_update_freq=None):
-
+                 predictor=None, reference_update_freq=None):
         """
         Args:
             equation_set (:class:`PrognosticEquationSet`): the prognostic
@@ -88,6 +89,15 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             accelerator (bool, optional): Whether to zero non-wind implicit
                 forcings for transport terms in order to speed up solver
                 convergence. Defaults to False.
+            predictor (str, optional): a single string corresponding to the name
+                of a variable to transport using the divergence predictor. This
+                pre-multiplies that variable by (1 - beta*dt*div(u)) before the
+                transport step, and calculates its transport increment from the
+                transport of this variable. This can improve the stability of
+                the time stepper at large time steps, when not using an
+                advective-then-flux formulation. This is only suitable for the
+                use on the conservative variable (e.g. depth or density).
+                Defaults to None, in which case no predictor is used.
             reference_update_freq (float, optional): frequency with which to
                 update the reference profile with the n-th time level state
                 fields. This variable corresponds to time in seconds, and
@@ -100,6 +110,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
         self.num_outer = num_outer
         self.num_inner = num_inner
         self.alpha = alpha
+        self.predictor = predictor
         self.accelerator = accelerator
         self.reference_update_freq = reference_update_freq
         self.to_update_ref_profile = False
@@ -201,6 +212,14 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
         self.forcing = Forcing(equation_set, self.alpha)
         self.bcs = equation_set.bcs
 
+        if self.predictor is not None:
+            V_DG = equation_set.domain.spaces('DG')
+            self.predictor_field_in = Function(V_DG)
+            div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.n('u'))
+            self.predictor_interpolator = Interpolator(
+                self.x.star(predictor)*div_factor, self.predictor_field_in
+            )
+
     def _apply_bcs(self):
         """
         Set the zero boundary conditions in the velocity.
@@ -263,6 +282,33 @@ def copy_active_tracers(self, x_in, x_out):
         for name in self.tracers_to_copy:
             x_out(name).assign(x_in(name))
 
+    def transport_field(self, name, scheme, xstar, xp):
+        """
+        Performs the transport of a field in xstar, placing the result in xp.
+
+        Args:
+            name (str): the name of the field to be transported.
+            scheme (:class:`TimeDiscretisation`): the time discretisation used
+                for the transport.
+            xstar (:class:`Fields`): the collection of state fields to be
+                transported.
+            xp (:class:`Fields`): the collection of state fields resulting from
+                the transport.
+        """
+
+        if name == self.predictor:
+            # Pre-multiply this variable by (1 - dt*beta*div(u))
+            V = xstar(name).function_space()
+            field_out = Function(V)
+            self.predictor_interpolator.interpolate()
+            scheme.apply(field_out, self.predictor_field_in)
+
+            # xp is xstar plus the increment from the transported predictor
+            xp(name).assign(xstar(name) + field_out - self.predictor_field_in)
+        else:
+            # Standard transport
+            scheme.apply(xp(name), xstar(name))
+
     def update_reference_profiles(self):
         """
         Updates the reference profiles and if required also updates them in the
@@ -324,7 +370,7 @@ def timestep(self):
                 for name, scheme in self.active_transport:
                     logger.info(f'Semi-implicit Quasi Newton: Transport {outer}: {name}')
                     # transports a field from xstar and puts result in xp
-                    scheme.apply(xp(name), xstar(name))
+                    self.transport_field(name, scheme, xstar, xp)
 
             # Fast physics -----------------------------------------------------
             x_after_fast(self.field_name).assign(xp(self.field_name))