PR #486: from firedrakeproject/SIQN_improvements

Improvements to the SIQN timeloop to make it closer to GungHo's

gusto/timeloop.py

@@ -479,7 +479,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
                  auxiliary_equations_and_schemes=None, linear_solver=None,
                  diffusion_schemes=None, physics_schemes=None,
                  slow_physics_schemes=None, fast_physics_schemes=None,
-                 alpha=Constant(0.5), num_outer=4, num_inner=1):
+                 alpha=Constant(0.5), off_centred_u=False,
+                 num_outer=4, num_inner=1):
 
         """
         Args:
@@ -516,6 +517,9 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             alpha (`ufl.Constant`, optional): the semi-implicit off-centering
                 parameter. A value of 1 corresponds to fully implicit, while 0
                 corresponds to fully explicit. Defaults to Constant(0.5).
+            off_centred_u (bool, optional): option to offcentre the transporting
+                velocity. Defaults to False, in which case transporting velocity
+                is centred. If True offcentring uses value of alpha.
             num_outer (int, optional): number of outer iterations in the semi-
                 implicit algorithm. The outer loop includes transport and any
                 fast physics schemes. Defaults to 4. Note that default used by
@@ -531,6 +535,10 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
         self.num_inner = num_inner
         self.alpha = alpha
 
+        # default is to not offcentre transporting velocity but if it
+        # is offcentred then use the same value as alpha
+        self.alpha_u = Constant(alpha) if off_centred_u else Constant(0.5)
+
         self.spatial_methods = spatial_methods
 
         if physics_schemes is not None:
@@ -637,7 +645,7 @@ def transporting_velocity(self):
         xn = self.x.n
         xnp1 = self.x.np1
         # computes ubar from un and unp1
-        return xn('u') + self.alpha*(xnp1('u')-xn('u'))
+        return xn('u') + self.alpha_u*(xnp1('u')-xn('u'))
 
     def setup_fields(self):
         """Sets up time levels n, star, p and np1"""
@@ -704,10 +712,11 @@ def timestep(self):
 
         with timed_stage("Apply forcing terms"):
             logger.info('SIQN: Explicit forcing')
-            # TODO: check if forcing is applied to x_after_slow or xn
             # Put explicit forcing into xstar
-            self.forcing.apply(x_after_slow, x_after_slow, xstar(self.field_name), "explicit")
+            self.forcing.apply(x_after_slow, xn, xstar(self.field_name), "explicit")
 
+        # set xp here so that variables that are not transported have
+        # the correct values
         xp(self.field_name).assign(xstar(self.field_name))
 
         for outer in range(self.num_outer):
@@ -732,9 +741,10 @@ def timestep(self):
 
             for inner in range(self.num_inner):
 
+                # TODO: this is where to update the reference state
+
                 with timed_stage("Apply forcing terms"):
                     logger.info(f'SIQN: Implicit forcing {(outer, inner)}')
-                    # TODO: why don't we update xnp1 with a better guess here?
                     self.forcing.apply(xp, xnp1, xrhs, "implicit")
 
                 xrhs -= xnp1(self.field_name)