From f516f6859bc9af14f28e5562420837a0b7bd7189 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Mon, 19 Aug 2024 09:00:30 +0100
Subject: [PATCH 1/9] allow reference profiles to be updated

---
 gusto/core/io.py                              | 49 ++++++++++++++++---
 gusto/solvers/linear_solvers.py               | 26 ++++++----
 .../semi_implicit_quasi_newton.py             | 44 ++++++++++++++---
 gusto/timestepping/timestepper.py             | 15 ++++--
 4 files changed, 105 insertions(+), 29 deletions(-)

diff --git a/gusto/core/io.py b/gusto/core/io.py
index a9a09f1fe..3c81ca4be 100644
--- a/gusto/core/io.py
+++ b/gusto/core/io.py
@@ -531,7 +531,11 @@ def setup_dump(self, state_fields, t, pick_up=False):
 
         # dump initial fields
         if not pick_up:
-            self.dump(state_fields, t, step=1)
+            step = 1
+            last_ref_update_time = None
+            initial_steps = None
+            time_data = (t, step, initial_steps, last_ref_update_time)
+            self.dump(state_fields, time_data)
 
     def pick_up_from_checkpoint(self, state_fields):
         """
@@ -541,7 +545,12 @@ def pick_up_from_checkpoint(self, state_fields):
             state_fields (:class:`StateFields`): the model's field container.
 
         Returns:
-            float: the checkpointed model time.
+            tuple of (`time_data`, `reference_profiles`): where `time_data`
+                itself is a tuple of numbers relating to the checkpointed time.
+                This tuple is: (model time, step index,
+                number of initial steps, last time reference profiles updated).
+                The `reference_profiles` are a list of (`field_name`, expr)
+                pairs describing the reference profile fields.
         """
 
         # -------------------------------------------------------------------- #
@@ -602,6 +611,13 @@ def pick_up_from_checkpoint(self, state_fields):
                     except AttributeError:
                         initial_steps = None
 
+                    # Try to pick up number last_ref_update_time
+                    # Not compulsory so errors allowed
+                    try:
+                        last_ref_update_time = chk.read_attribute("/", "last_ref_update_time")
+                    except AttributeError:
+                        last_ref_update_time = None
+
                     # Finally pick up time and step number
                     t = chk.read_attribute("/", "time")
                     step = chk.read_attribute("/", "step")
@@ -632,6 +648,13 @@ def pick_up_from_checkpoint(self, state_fields):
                     else:
                         initial_steps = None
 
+                    # Try to pick up last reference profile update time
+                    # Not compulsory so errors allowed
+                    if chk.has_attr("/", "last_ref_update_time"):
+                        last_ref_update_time = chk.get_attr("/", "last_ref_update_time")
+                    else:
+                        last_ref_update_time = None
+
                     # Finally pick up time
                     t = chk.get_attr("/", "time")
                     step = chk.get_attr("/", "step")
@@ -647,9 +670,10 @@ def pick_up_from_checkpoint(self, state_fields):
             if hasattr(diagnostic_field, "init_field_set"):
                 diagnostic_field.init_field_set = True
 
-        return t, reference_profiles, step, initial_steps
+        time_data = (t, step, initial_steps, last_ref_update_time)
+        return time_data, reference_profiles
 
-    def dump(self, state_fields, t, step, initial_steps=None):
+    def dump(self, state_fields, time_data):
         """
         Dumps all of the required model output.
 
@@ -659,12 +683,17 @@ def dump(self, state_fields, t, step, initial_steps=None):
 
         Args:
             state_fields (:class:`StateFields`): the model's field container.
-            t (float): the simulation's current time.
-            step (int): the number of time steps.
-            initial_steps (int, optional): the number of initial time steps
-                completed by a multi-level time scheme. Defaults to None.
+            time_data (tuple): contains information relating to the time in
+                the simulation. The tuple is structured as follows:
+                - t: current time in s
+                - step: the index of the time step
+                - initial_steps: number of initial time steps completed by a
+                  multi-level time scheme (could be None)
+                - last_ref_update_time: the last time in s that the reference
+                  profiles were updated (could be None)
         """
         output = self.output
+        t, step, initial_steps, last_ref_update_time = time_data
 
         # Diagnostics:
         # Compute diagnostic fields
@@ -688,6 +717,8 @@ def dump(self, state_fields, t, step, initial_steps=None):
                 self.chkpt.write_attribute("/", "step", step)
                 if initial_steps is not None:
                     self.chkpt.write_attribute("/", "initial_steps", initial_steps)
+                if last_ref_update_time is not None:
+                    self.chkpt.write_attribute("/", "last_ref_update_time", last_ref_update_time)
             else:
                 with CheckpointFile(self.chkpt_path, 'w') as chk:
                     chk.save_mesh(self.domain.mesh)
@@ -697,6 +728,8 @@ def dump(self, state_fields, t, step, initial_steps=None):
                     chk.set_attr("/", "step", step)
                     if initial_steps is not None:
                         chk.set_attr("/", "initial_steps", initial_steps)
+                    if last_ref_update_time is not None:
+                        chk.set_attr("/", "last_ref_update_time", last_ref_update_time)
 
         if (next(self.dumpcount) % output.dumpfreq) == 0:
             if output.dump_nc:
diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index f3a610da7..cfef6aca9 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -27,7 +27,8 @@
 from abc import ABCMeta, abstractmethod, abstractproperty
 
 
-__all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver", "MoistConvectiveSWSolver"]
+__all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver",
+           "ThermalSWSolver", "MoistConvectiveSWSolver"]
 
 
 class TimesteppingSolver(object, metaclass=ABCMeta):
@@ -371,6 +372,20 @@ def L_tr(f):
         python_context = self.hybridized_solver.snes.ksp.pc.getPythonContext()
         attach_custom_monitor(python_context, logging_ksp_monitor_true_residual)
 
+    @timed_function("Gusto:UpdateReferenceProfiles")
+    def update_reference_profiles(self):
+        """
+        Updates the reference profiles.
+        """
+
+        with timed_region("Gusto:HybridProjectRhobar"):
+            logger.info('Compressible linear solver: rho average solve')
+            self.rho_avg_solver.solve()
+
+        with timed_region("Gusto:HybridProjectExnerbar"):
+            logger.info('Compressible linear solver: Exner average solve')
+            self.exner_avg_solver.solve()
+
     @timed_function("Gusto:LinearSolve")
     def solve(self, xrhs, dy):
         """
@@ -384,15 +399,6 @@ def solve(self, xrhs, dy):
         """
         self.xrhs.assign(xrhs)
 
-        # TODO: can we avoid computing these each time the solver is called?
-        with timed_region("Gusto:HybridProjectRhobar"):
-            logger.info('Compressible linear solver: rho average solve')
-            self.rho_avg_solver.solve()
-
-        with timed_region("Gusto:HybridProjectExnerbar"):
-            logger.info('Compressible linear solver: Exner average solve')
-            self.exner_avg_solver.solve()
-
         # Solve the hybridized system
         logger.info('Compressible linear solver: hybridized solve')
         self.hybridized_solver.solve()
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index 1e524a100..2bed298d0 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -35,7 +35,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
                  diffusion_schemes=None, physics_schemes=None,
                  slow_physics_schemes=None, fast_physics_schemes=None,
                  alpha=Constant(0.5), off_centred_u=False,
-                 num_outer=2, num_inner=2, accelerator=False):
+                 num_outer=2, num_inner=2, accelerator=False,
+                 reference_update_freq=None):
 
         """
         Args:
@@ -84,13 +85,23 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
                 implicit forcing (pressure gradient and Coriolis) terms, and the
                 linear solve. Defaults to 2. Note that default used by the Met
                 Office's ENDGame and GungHo models is 2.
-            accelerator (bool, optional): Whether to zero non-wind implicit forcings
-                for transport terms in order to speed up solver convergence
+            accelerator (bool, optional): Whether to zero non-wind implicit
+                forcings for transport terms in order to speed up solver
+                convergence. Defaults to False.
+            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
+                setting this to zero will update the reference profiles every
+                time step. Setting it to None turns off the update, and
+                reference profiles will remain at their initial values.
+                Defaults to None.
         """
 
         self.num_outer = num_outer
         self.num_inner = num_inner
         self.alpha = alpha
+        self.accelerator = accelerator
+        self.reference_update_freq = reference_update_freq
 
         # default is to not offcentre transporting velocity but if it
         # is offcentred then use the same value as alpha
@@ -188,7 +199,6 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             self.linear_solver = linear_solver
         self.forcing = Forcing(equation_set, self.alpha)
         self.bcs = equation_set.bcs
-        self.accelerator = accelerator
 
     def _apply_bcs(self):
         """
@@ -264,6 +274,15 @@ def timestep(self):
         xrhs_phys = self.xrhs_phys
         dy = self.dy
 
+        # Update reference profiles --------------------------------------------
+        if self.reference_update_freq is not None:
+            if float(self.t) + self.reference_update_freq > self.last_ref_update_time:
+                self.equation.X_ref.assign(self.x.n)
+                self.last_ref_update_time = float(self.t)
+                if hasattr(self.linear_solver, 'update_reference_profiles'):
+                    self.linear_solver.update_reference_profiles()
+
+        # Slow physics ---------------------------------------------------------
         x_after_slow(self.field_name).assign(xn(self.field_name))
         if len(self.slow_physics_schemes) > 0:
             with timed_stage("Slow physics"):
@@ -271,6 +290,7 @@ def timestep(self):
                 for _, scheme in self.slow_physics_schemes:
                     scheme.apply(x_after_slow(scheme.field_name), x_after_slow(scheme.field_name))
 
+        # Explict forcing ------------------------------------------------------
         with timed_stage("Apply forcing terms"):
             logger.info('Semi-implicit Quasi Newton: Explicit forcing')
             # Put explicit forcing into xstar
@@ -280,8 +300,10 @@ def timestep(self):
         # the correct values
         xp(self.field_name).assign(xstar(self.field_name))
 
+        # OUTER ----------------------------------------------------------------
         for outer in range(self.num_outer):
 
+            # Transport --------------------------------------------------------
             with timed_stage("Transport"):
                 self.io.log_courant(self.fields, 'transporting_velocity',
                                     message=f'transporting velocity, outer iteration {outer}')
@@ -290,6 +312,7 @@ def timestep(self):
                     # transports a field from xstar and puts result in xp
                     scheme.apply(xp(name), xstar(name))
 
+            # Fast physics -----------------------------------------------------
             x_after_fast(self.field_name).assign(xp(self.field_name))
             if len(self.fast_physics_schemes) > 0:
                 with timed_stage("Fast physics"):
@@ -302,8 +325,7 @@ def timestep(self):
 
             for inner in range(self.num_inner):
 
-                # TODO: this is where to update the reference state
-
+                # Implicit forcing ---------------------------------------------
                 with timed_stage("Apply forcing terms"):
                     logger.info(f'Semi-implicit Quasi Newton: Implicit forcing {(outer, inner)}')
                     self.forcing.apply(xp, xnp1, xrhs, "implicit")
@@ -314,6 +336,7 @@ def timestep(self):
                 xrhs -= xnp1(self.field_name)
                 xrhs += xrhs_phys
 
+                # Linear solve -------------------------------------------------
                 with timed_stage("Implicit solve"):
                     logger.info(f'Semi-implicit Quasi Newton: Mixed solve {(outer, inner)}')
                     self.linear_solver.solve(xrhs, dy)  # solves linear system and places result in dy
@@ -353,10 +376,17 @@ def run(self, t, tmax, pick_up=False):
             pick_up: (bool): specify whether to pick_up from a previous run
         """
 
-        if not pick_up:
+        if not pick_up and self.reference_update_freq is None:
             assert self.reference_profiles_initialised, \
                 'Reference profiles for must be initialised to use Semi-Implicit Timestepper'
 
+        if not pick_up and self.reference_update_freq is not None:
+            # Force reference profiles to be updated on first time step
+            self.last_ref_update_time = t - float(self.dt)
+        elif not pick_up:
+            if hasattr(self.linear_solver, 'update_reference_profiles'):
+                self.linear_solver.update_reference_profiles()
+
         super().run(t, tmax, pick_up=pick_up)
 
 
diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py
index d4f63be8a..38c8323bb 100644
--- a/gusto/timestepping/timestepper.py
+++ b/gusto/timestepping/timestepper.py
@@ -30,6 +30,7 @@ def __init__(self, equation, io):
         self.dt = self.equation.domain.dt
         self.t = self.equation.domain.t
         self.reference_profiles_initialised = False
+        self.last_ref_update_time = None
 
         self.setup_fields()
         self.setup_scheme()
@@ -163,8 +164,9 @@ def run(self, t, tmax, pick_up=False):
 
         if pick_up:
             # Pick up fields, and return other info to be picked up
-            t, reference_profiles, self.step, initial_timesteps = self.io.pick_up_from_checkpoint(self.fields)
-            self.set_reference_profiles(reference_profiles)
+            time_data, reference_profiles = self.io.pick_up_from_checkpoint(self.fields)
+            t, self.step, initial_timesteps, last_ref_update_time = time_data
+            self.set_reference_profiles(reference_profiles, last_ref_update_time)
             self.set_initial_timesteps(initial_timesteps)
         else:
             self.step = 1
@@ -193,14 +195,15 @@ def run(self, t, tmax, pick_up=False):
             self.step += 1
 
             with timed_stage("Dump output"):
-                self.io.dump(self.fields, float(self.t), self.step, self.get_initial_timesteps())
+                time_data = (float(t), self.step, self.get_initial_timesteps(), self.last_ref_update_time)
+                self.io.dump(self.fields, time_data)
 
         if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint':
             self.io.chkpt.close()
 
         logger.info(f'TIMELOOP complete. t={float(self.t):.5f}, {tmax=:.5f}')
 
-    def set_reference_profiles(self, reference_profiles):
+    def set_reference_profiles(self, reference_profiles, last_ref_update_time=None):
         """
         Initialise the model's reference profiles.
 
@@ -208,6 +211,8 @@ def set_reference_profiles(self, reference_profiles):
             where 'field_name' is the string giving the name of the reference
             profile field expr is the :class:`ufl.Expr` whose value is used to
             set the reference field.
+        last_ref_update_time (float, optional): the last time that the reference
+            profiles were updated. Defaults to None.
         """
         for field_name, profile in reference_profiles:
             if field_name+'_bar' in self.fields:
@@ -237,6 +242,8 @@ def set_reference_profiles(self, reference_profiles):
                     # Don't need to do anything else as value in field container has already been set
         self.reference_profiles_initialised = True
 
+        self.last_ref_update_time = last_ref_update_time
+
 
 class Timestepper(BaseTimestepper):
     """

From cd133ba74109768bba95d37536b983a8b038ac25 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Mon, 19 Aug 2024 21:16:53 +0100
Subject: [PATCH 2/9] get updating reference profile working with the help of
 the checkpointing test

---
 .../semi_implicit_quasi_newton.py             | 35 +++++++++++++------
 gusto/timestepping/timestepper.py             |  6 +++-
 integration-tests/model/test_checkpointing.py | 18 +++++-----
 3 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index 2bed298d0..ce8758943 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -102,6 +102,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
         self.alpha = alpha
         self.accelerator = accelerator
         self.reference_update_freq = reference_update_freq
+        self.to_update_ref_profile = False
 
         # default is to not offcentre transporting velocity but if it
         # is offcentred then use the same value as alpha
@@ -262,6 +263,24 @@ def copy_active_tracers(self, x_in, x_out):
         for name in self.tracers_to_copy:
             x_out(name).assign(x_in(name))
 
+    def update_reference_profiles(self):
+        """
+        Updates the reference profiles and if required also updates them in the
+        linear solver.
+        """
+
+        if self.reference_update_freq is not None:
+            if float(self.t) + self.reference_update_freq > self.last_ref_update_time:
+                self.equation.X_ref.assign(self.x.n(self.field_name))
+                self.last_ref_update_time = float(self.t)
+                if hasattr(self.linear_solver, 'update_reference_profiles'):
+                    self.linear_solver.update_reference_profiles()
+
+        elif self.to_update_ref_profile:
+            if hasattr(self.linear_solver, 'update_reference_profiles'):
+                self.linear_solver.update_reference_profiles()
+                self.to_update_ref_profile = False
+
     def timestep(self):
         """Defines the timestep"""
         xn = self.x.n
@@ -275,12 +294,7 @@ def timestep(self):
         dy = self.dy
 
         # Update reference profiles --------------------------------------------
-        if self.reference_update_freq is not None:
-            if float(self.t) + self.reference_update_freq > self.last_ref_update_time:
-                self.equation.X_ref.assign(self.x.n)
-                self.last_ref_update_time = float(self.t)
-                if hasattr(self.linear_solver, 'update_reference_profiles'):
-                    self.linear_solver.update_reference_profiles()
+        self.update_reference_profiles()
 
         # Slow physics ---------------------------------------------------------
         x_after_slow(self.field_name).assign(xn(self.field_name))
@@ -382,10 +396,11 @@ def run(self, t, tmax, pick_up=False):
 
         if not pick_up and self.reference_update_freq is not None:
             # Force reference profiles to be updated on first time step
-            self.last_ref_update_time = t - float(self.dt)
-        elif not pick_up:
-            if hasattr(self.linear_solver, 'update_reference_profiles'):
-                self.linear_solver.update_reference_profiles()
+            self.last_ref_update_time = float(t) - float(self.dt)
+
+        elif not pick_up or (pick_up and self.reference_update_freq is None):
+            # Indicate that linear solver profile needs updating
+            self.to_update_ref_profile = True
 
         super().run(t, tmax, pick_up=pick_up)
 
diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py
index 38c8323bb..45eb603ab 100644
--- a/gusto/timestepping/timestepper.py
+++ b/gusto/timestepping/timestepper.py
@@ -168,6 +168,7 @@ def run(self, t, tmax, pick_up=False):
             t, self.step, initial_timesteps, last_ref_update_time = time_data
             self.set_reference_profiles(reference_profiles, last_ref_update_time)
             self.set_initial_timesteps(initial_timesteps)
+
         else:
             self.step = 1
 
@@ -195,7 +196,10 @@ def run(self, t, tmax, pick_up=False):
             self.step += 1
 
             with timed_stage("Dump output"):
-                time_data = (float(t), self.step, self.get_initial_timesteps(), self.last_ref_update_time)
+                time_data = (
+                    float(self.t), self.step,
+                    self.get_initial_timesteps(), self.last_ref_update_time
+                )
                 self.io.dump(self.fields, time_data)
 
         if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint':
diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py
index 28549a66f..dc72fb1de 100644
--- a/integration-tests/model/test_checkpointing.py
+++ b/integration-tests/model/test_checkpointing.py
@@ -11,7 +11,7 @@
 import pytest
 
 
-def set_up_model_objects(mesh, dt, output, stepper_type):
+def set_up_model_objects(mesh, dt, output, stepper_type, ref_update_freq):
 
     domain = Domain(mesh, dt, "CG", 1)
 
@@ -40,7 +40,8 @@ def set_up_model_objects(mesh, dt, output, stepper_type):
         # build time stepper
         stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields,
                                           transport_methods,
-                                          linear_solver=linear_solver)
+                                          linear_solver=linear_solver,
+                                          reference_update_freq=ref_update_freq)
 
     elif stepper_type == 'multi_level':
         scheme = AdamsBashforth(domain, order=2)
@@ -92,9 +93,10 @@ def initialise_fields(eqns, stepper):
     stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
 
 
-@pytest.mark.parametrize("stepper_type", ["multi_level", "semi_implicit"])
+@pytest.mark.parametrize("stepper_type, ref_update_freq", [
+    ("multi_level", None), ("semi_implicit", None), ("semi_implicit", 0.6)])
 @pytest.mark.parametrize("checkpoint_method", ["dumbcheckpoint", "checkpointfile"])
-def test_checkpointing(tmpdir, stepper_type, checkpoint_method):
+def test_checkpointing(tmpdir, stepper_type, checkpoint_method, ref_update_freq):
 
     mesh_name = 'checkpointing_mesh'
 
@@ -128,8 +130,8 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method):
         chkptfreq=2,
     )
 
-    stepper_1, eqns_1 = set_up_model_objects(mesh, dt, output_1, stepper_type)
-    stepper_2, eqns_2 = set_up_model_objects(mesh, dt, output_2, stepper_type)
+    stepper_1, eqns_1 = set_up_model_objects(mesh, dt, output_1, stepper_type, ref_update_freq)
+    stepper_2, eqns_2 = set_up_model_objects(mesh, dt, output_2, stepper_type, ref_update_freq)
 
     initialise_fields(eqns_1, stepper_1)
     initialise_fields(eqns_2, stepper_2)
@@ -163,7 +165,7 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method):
 
     if checkpoint_method == 'checkpointfile':
         mesh = pick_up_mesh(output_3, mesh_name)
-    stepper_3, _ = set_up_model_objects(mesh, dt, output_3, stepper_type)
+    stepper_3, _ = set_up_model_objects(mesh, dt, output_3, stepper_type, ref_update_freq)
     stepper_3.io.pick_up_from_checkpoint(stepper_3.fields)
 
     # ------------------------------------------------------------------------ #
@@ -192,7 +194,7 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method):
     )
     if checkpoint_method == 'checkpointfile':
         mesh = pick_up_mesh(output_3, mesh_name)
-    stepper_3, _ = set_up_model_objects(mesh, dt, output_3, stepper_type)
+    stepper_3, _ = set_up_model_objects(mesh, dt, output_3, stepper_type, ref_update_freq)
     stepper_3.run(t=2*dt, tmax=4*dt, pick_up=True)
 
     # ------------------------------------------------------------------------ #

From 3e75c5b1bbef58424dd0a3973fa5bc6004214f6b Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Thu, 12 Sep 2024 17:47:23 +0100
Subject: [PATCH 3/9] investigate subcycling

---
 .../skamarock_klemp_nonhydrostatic.py              | 13 +++++++------
 examples/shallow_water/williamson_5.py             | 14 ++++++++++----
 gusto/time_discretisation/time_discretisation.py   |  1 +
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
index 777134d61..65ec7c831 100644
--- a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
+++ b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
@@ -83,7 +83,7 @@ def skamarock_klemp_nonhydrostatic(
     if COMM_WORLD.size == 1:
         output = OutputParameters(
             dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq,
-            dump_vtus=True, dump_nc=False,
+            dump_vtus=False, dump_nc=True,
             point_data=[('theta_perturbation', points)],
         )
     else:
@@ -106,9 +106,9 @@ def skamarock_klemp_nonhydrostatic(
     # Transport schemes
     theta_opts = SUPGOptions()
     transported_fields = [
-        TrapeziumRule(domain, "u"),
-        SSPRK3(domain, "rho"),
-        SSPRK3(domain, "theta", options=theta_opts)
+        SSPRK3(domain, "u", subcycle_by_courant=0.2),
+        SSPRK3(domain, "rho", subcycle_by_courant=0.2),
+        SSPRK3(domain, "theta", subcycle_by_courant=0.2, options=theta_opts)
     ]
     transport_methods = [
         DGUpwind(eqns, "u"),
@@ -117,12 +117,13 @@ def skamarock_klemp_nonhydrostatic(
     ]
 
     # Linear solver
-    linear_solver = CompressibleSolver(eqns)
+    linear_solver = CompressibleSolver(eqns, tau_values={'theta': 1.0, 'rho': 1.0})
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
         eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver
+        linear_solver=linear_solver, alpha=0.5, reference_update_freq=0.0,
+        num_outer=4, num_inner=1, accelerator=True
     )
 
     # ------------------------------------------------------------------------ #
diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py
index cdb2b58a9..50bb9b379 100644
--- a/examples/shallow_water/williamson_5.py
+++ b/examples/shallow_water/williamson_5.py
@@ -14,7 +14,7 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, Sum,
     lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent,
-    MeridionalComponent, RelativeVorticity
+    MeridionalComponent, RelativeVorticity, MoistConvectiveSWSolver
 )
 
 williamson_5_defaults = {
@@ -76,19 +76,25 @@ def williamson_5(
     # I/O
     output = OutputParameters(
         dirname=dirname, dumplist_latlon=['D'], dumpfreq=dumpfreq,
-        dump_vtus=True, dump_nc=False, dumplist=['D', 'topography']
+        dump_vtus=False, dump_nc=True, dumplist=['D', 'topography']
     )
     diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(),
                          MeridionalComponent('u'), ZonalComponent('u')]
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
-    transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")]
+    transported_fields = [
+        SSPRK3(domain, "u", subcycle_by_courant=0.25),
+        SSPRK3(domain, "D", subcycle_by_courant=0.25)
+    ]
     transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
 
+    linear_solver = MoistConvectiveSWSolver(eqns, tau_values={'D': 1.0})
+
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods
+        eqns, io, transported_fields, transport_methods,
+        linear_solver=linear_solver, num_outer=4, num_inner=1
     )
 
     # ------------------------------------------------------------------------ #
diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py
index c0438295c..2e49e77bc 100644
--- a/gusto/time_discretisation/time_discretisation.py
+++ b/gusto/time_discretisation/time_discretisation.py
@@ -420,6 +420,7 @@ def apply(self, x_out, x_in):
 
         self.x0.assign(x_in)
         for i in range(self.ncycles):
+            logger.info(f'Explicit time discretisation, subcycle {i}, dt = {float(self.dt)}')
             self.apply_cycle(self.x1, self.x0)
             self.x0.assign(self.x1)
         x_out.assign(self.x1)

From ad978385ad09e3fbcd7a0529325e761c3a0a9dcc Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Mon, 16 Sep 2024 08:30:47 +0100
Subject: [PATCH 4/9] implement predictors and explore some settings

---
 .../skamarock_klemp_nonhydrostatic.py         |  4 ++--
 examples/shallow_water/williamson_5.py        |  6 +++--
 .../time_discretisation.py                    |  2 +-
 .../semi_implicit_quasi_newton.py             | 23 ++++++++++++++++---
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
index 65ec7c831..56087c789 100644
--- a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
+++ b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
@@ -122,8 +122,8 @@ def skamarock_klemp_nonhydrostatic(
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
         eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver, alpha=0.5, reference_update_freq=0.0,
-        num_outer=4, num_inner=1, accelerator=True
+        linear_solver=linear_solver, alpha=0.55, predictor='rho',
+        num_outer=2, num_inner=2, accelerator=True
     )
 
     # ------------------------------------------------------------------------ #
diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py
index 50bb9b379..f1e9afd2d 100644
--- a/examples/shallow_water/williamson_5.py
+++ b/examples/shallow_water/williamson_5.py
@@ -71,7 +71,8 @@ def williamson_5(
     rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2)
     r = sqrt(rsq)
     tpexpr = mountain_height * (1 - r/R0)
-    eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr)
+    eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr,
+                                 u_transport_option='vector_advection_form')
 
     # I/O
     output = OutputParameters(
@@ -94,7 +95,8 @@ def williamson_5(
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
         eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver, num_outer=4, num_inner=1
+        linear_solver=linear_solver, num_outer=2, num_inner=2,
+        predictor='D', alpha=0.55, accelerator=True
     )
 
     # ------------------------------------------------------------------------ #
diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py
index 2e49e77bc..38d94f073 100644
--- a/gusto/time_discretisation/time_discretisation.py
+++ b/gusto/time_discretisation/time_discretisation.py
@@ -415,7 +415,7 @@ def apply(self, x_out, x_in):
         """
         # If doing adaptive subcycles, update dt and ncycles here
         if self.subcycle_by_courant is not None:
-            self.ncycles = math.ceil(float(self.courant_max)/self.subcycle_by_courant)
+            self.ncycles = min(math.ceil(float(self.courant_max)/self.subcycle_by_courant), 20)
             self.dt.assign(self.original_dt/self.ncycles)
 
         self.x0.assign(x_in)
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index ce8758943..a835d7b66 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -4,7 +4,8 @@
 """
 
 from firedrake import (Function, Constant, TrialFunctions, DirichletBC,
-                       LinearVariationalProblem, LinearVariationalSolver)
+                       LinearVariationalProblem, LinearVariationalSolver,
+                       Interpolator, div)
 from firedrake.fml import drop, replace_subject
 from pyop2.profiling import timed_stage
 from gusto.core import TimeLevelFields, StateFields
@@ -36,7 +37,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):
+                 reference_update_freq=None, predictor=None):
 
         """
         Args:
@@ -103,6 +104,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
         self.accelerator = accelerator
         self.reference_update_freq = reference_update_freq
         self.to_update_ref_profile = False
+        self.predictor = predictor
 
         # default is to not offcentre transporting velocity but if it
         # is offcentred then use the same value as alpha
@@ -201,6 +203,13 @@ 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')
+            div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.star('u'))
+            self.predictor_interpolator = Interpolator(
+                self.x.star(predictor)*div_factor, V_DG
+            )
+
     def _apply_bcs(self):
         """
         Set the zero boundary conditions in the velocity.
@@ -324,7 +333,15 @@ 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))
+                    if name == self.predictor:
+                        V = xstar(name).function_space()
+                        field_in = Function(V)
+                        field_out = Function(V)
+                        self.predictor_interpolator.interpolate()
+                        scheme.apply(field_out, field_in)
+                        xp(name).assign(xstar(name) + field_out - field_in)
+                    else:
+                        scheme.apply(xp(name), xstar(name))
 
             # Fast physics -----------------------------------------------------
             x_after_fast(self.field_name).assign(xp(self.field_name))

From f06e8a86056180d2be8957a792f08feffcfbe52d Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Sat, 21 Sep 2024 13:10:55 +0100
Subject: [PATCH 5/9] tidy up example

---
 examples/shallow_water/williamson_5.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py
index f1e9afd2d..5e46f730a 100644
--- a/examples/shallow_water/williamson_5.py
+++ b/examples/shallow_water/williamson_5.py
@@ -14,7 +14,7 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, Sum,
     lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent,
-    MeridionalComponent, RelativeVorticity, MoistConvectiveSWSolver
+    MeridionalComponent, RelativeVorticity,
 )
 
 williamson_5_defaults = {
@@ -52,6 +52,7 @@ def williamson_5(
     # ------------------------------------------------------------------------ #
 
     element_order = 1
+
     # ------------------------------------------------------------------------ #
     # Set up model objects
     # ------------------------------------------------------------------------ #
@@ -85,18 +86,14 @@ def williamson_5(
 
     # Transport schemes
     transported_fields = [
-        SSPRK3(domain, "u", subcycle_by_courant=0.25),
-        SSPRK3(domain, "D", subcycle_by_courant=0.25)
+        TrapeziumRule(domain, "u"),
+        SSPRK3(domain, "D", subcycle_by_courant=0.3)
     ]
     transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
 
-    linear_solver = MoistConvectiveSWSolver(eqns, tau_values={'D': 1.0})
-
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver, num_outer=2, num_inner=2,
-        predictor='D', alpha=0.55, accelerator=True
+        eqns, io, transported_fields, transport_methods, predictor='D'
     )
 
     # ------------------------------------------------------------------------ #

From 0678bae35ef00249f1fdf6f8bfddbae0afe11e6b Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Sat, 21 Sep 2024 13:29:06 +0100
Subject: [PATCH 6/9] tidy up and add comment

---
 examples/shallow_water/williamson_5.py        |  7 +-
 .../semi_implicit_quasi_newton.py             | 99 ++++++++-----------
 2 files changed, 46 insertions(+), 60 deletions(-)

diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py
index 5e46f730a..575eeef08 100644
--- a/examples/shallow_water/williamson_5.py
+++ b/examples/shallow_water/williamson_5.py
@@ -14,7 +14,7 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, Sum,
     lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent,
-    MeridionalComponent, RelativeVorticity,
+    MeridionalComponent, RelativeVorticity
 )
 
 williamson_5_defaults = {
@@ -72,13 +72,12 @@ def williamson_5(
     rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2)
     r = sqrt(rsq)
     tpexpr = mountain_height * (1 - r/R0)
-    eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr,
-                                 u_transport_option='vector_advection_form')
+    eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr)
 
     # I/O
     output = OutputParameters(
         dirname=dirname, dumplist_latlon=['D'], dumpfreq=dumpfreq,
-        dump_vtus=False, dump_nc=True, dumplist=['D', 'topography']
+        dump_vtus=True, dump_nc=False, dumplist=['D', 'topography']
     )
     diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(),
                          MeridionalComponent('u'), ZonalComponent('u')]
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index a835d7b66..3a645046c 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -3,9 +3,10 @@
 and GungHo dynamical cores.
 """
 
-from firedrake import (Function, Constant, TrialFunctions, DirichletBC,
-                       LinearVariationalProblem, LinearVariationalSolver,
-                       Interpolator, div)
+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
@@ -37,7 +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):
+                 predictor=None):
 
         """
         Args:
@@ -89,21 +90,20 @@ 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.
-            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
-                setting this to zero will update the reference profiles every
-                time step. Setting it to None turns off the update, and
-                reference profiles will remain at their initial values.
-                Defaults to None.
+            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.
         """
 
         self.num_outer = num_outer
         self.num_inner = num_inner
         self.alpha = alpha
-        self.accelerator = accelerator
-        self.reference_update_freq = reference_update_freq
-        self.to_update_ref_profile = False
         self.predictor = predictor
 
         # default is to not offcentre transporting velocity but if it
@@ -202,6 +202,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             self.linear_solver = linear_solver
         self.forcing = Forcing(equation_set, self.alpha)
         self.bcs = equation_set.bcs
+        self.accelerator = accelerator
 
         if self.predictor is not None:
             V_DG = equation_set.domain.spaces('DG')
@@ -272,23 +273,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 update_reference_profiles(self):
+    def transport_field(self, name, scheme, xstar, xp):
         """
-        Updates the reference profiles and if required also updates them in the
-        linear solver.
+        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 self.reference_update_freq is not None:
-            if float(self.t) + self.reference_update_freq > self.last_ref_update_time:
-                self.equation.X_ref.assign(self.x.n(self.field_name))
-                self.last_ref_update_time = float(self.t)
-                if hasattr(self.linear_solver, 'update_reference_profiles'):
-                    self.linear_solver.update_reference_profiles()
+        if name == self.predictor:
+            # Pre-multiply this variable by (1 - dt*beta*div(u))
+            V = xstar(name).function_space()
+            field_in = Function(V)
+            field_out = Function(V)
+            self.predictor_interpolator.interpolate()
+            scheme.apply(field_out, field_in)
 
-        elif self.to_update_ref_profile:
-            if hasattr(self.linear_solver, 'update_reference_profiles'):
-                self.linear_solver.update_reference_profiles()
-                self.to_update_ref_profile = False
+            # xp is xstar plus the increment from the transported predictor
+            xp(name).assign(xstar(name) + field_out - field_in)
+        else:
+            # Standard transport
+            scheme.apply(xp(name), xstar(name))
 
     def timestep(self):
         """Defines the timestep"""
@@ -302,10 +313,6 @@ def timestep(self):
         xrhs_phys = self.xrhs_phys
         dy = self.dy
 
-        # Update reference profiles --------------------------------------------
-        self.update_reference_profiles()
-
-        # Slow physics ---------------------------------------------------------
         x_after_slow(self.field_name).assign(xn(self.field_name))
         if len(self.slow_physics_schemes) > 0:
             with timed_stage("Slow physics"):
@@ -313,7 +320,6 @@ def timestep(self):
                 for _, scheme in self.slow_physics_schemes:
                     scheme.apply(x_after_slow(scheme.field_name), x_after_slow(scheme.field_name))
 
-        # Explict forcing ------------------------------------------------------
         with timed_stage("Apply forcing terms"):
             logger.info('Semi-implicit Quasi Newton: Explicit forcing')
             # Put explicit forcing into xstar
@@ -323,27 +329,16 @@ def timestep(self):
         # the correct values
         xp(self.field_name).assign(xstar(self.field_name))
 
-        # OUTER ----------------------------------------------------------------
         for outer in range(self.num_outer):
 
-            # Transport --------------------------------------------------------
             with timed_stage("Transport"):
                 self.io.log_courant(self.fields, 'transporting_velocity',
                                     message=f'transporting velocity, outer iteration {outer}')
                 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
-                    if name == self.predictor:
-                        V = xstar(name).function_space()
-                        field_in = Function(V)
-                        field_out = Function(V)
-                        self.predictor_interpolator.interpolate()
-                        scheme.apply(field_out, field_in)
-                        xp(name).assign(xstar(name) + field_out - field_in)
-                    else:
-                        scheme.apply(xp(name), xstar(name))
-
-            # Fast physics -----------------------------------------------------
+                    self.transport_field(name, scheme, xstar, xp)
+
             x_after_fast(self.field_name).assign(xp(self.field_name))
             if len(self.fast_physics_schemes) > 0:
                 with timed_stage("Fast physics"):
@@ -356,7 +351,8 @@ def timestep(self):
 
             for inner in range(self.num_inner):
 
-                # Implicit forcing ---------------------------------------------
+                # TODO: this is where to update the reference state
+
                 with timed_stage("Apply forcing terms"):
                     logger.info(f'Semi-implicit Quasi Newton: Implicit forcing {(outer, inner)}')
                     self.forcing.apply(xp, xnp1, xrhs, "implicit")
@@ -367,7 +363,6 @@ def timestep(self):
                 xrhs -= xnp1(self.field_name)
                 xrhs += xrhs_phys
 
-                # Linear solve -------------------------------------------------
                 with timed_stage("Implicit solve"):
                     logger.info(f'Semi-implicit Quasi Newton: Mixed solve {(outer, inner)}')
                     self.linear_solver.solve(xrhs, dy)  # solves linear system and places result in dy
@@ -407,18 +402,10 @@ def run(self, t, tmax, pick_up=False):
             pick_up: (bool): specify whether to pick_up from a previous run
         """
 
-        if not pick_up and self.reference_update_freq is None:
+        if not pick_up:
             assert self.reference_profiles_initialised, \
                 'Reference profiles for must be initialised to use Semi-Implicit Timestepper'
 
-        if not pick_up and self.reference_update_freq is not None:
-            # Force reference profiles to be updated on first time step
-            self.last_ref_update_time = float(t) - float(self.dt)
-
-        elif not pick_up or (pick_up and self.reference_update_freq is None):
-            # Indicate that linear solver profile needs updating
-            self.to_update_ref_profile = True
-
         super().run(t, tmax, pick_up=pick_up)
 
 

From bf8811dd59d609ada83cb3c22dc8ccdb2040a991 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Mon, 23 Sep 2024 08:47:28 +0100
Subject: [PATCH 7/9] use xn for the predictor instead of xstar (makes no
 difference but matches gungho)

---
 gusto/timestepping/semi_implicit_quasi_newton.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index 3a645046c..d8cb0ab1a 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -206,7 +206,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
 
         if self.predictor is not None:
             V_DG = equation_set.domain.spaces('DG')
-            div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.star('u'))
+            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, V_DG
             )

From 279ea8bf0cad8233b39e322945c25118d73ef4e0 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Thu, 26 Sep 2024 08:50:05 +0100
Subject: [PATCH 8/9] actually make predictor work...

---
 gusto/timestepping/semi_implicit_quasi_newton.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index d8cb0ab1a..75dad6db4 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -206,9 +206,10 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
 
         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, V_DG
+                self.x.star(predictor)*div_factor, self.predictor_field_in
             )
 
     def _apply_bcs(self):
@@ -290,13 +291,12 @@ def transport_field(self, name, scheme, xstar, xp):
         if name == self.predictor:
             # Pre-multiply this variable by (1 - dt*beta*div(u))
             V = xstar(name).function_space()
-            field_in = Function(V)
             field_out = Function(V)
             self.predictor_interpolator.interpolate()
-            scheme.apply(field_out, field_in)
+            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 - field_in)
+            xp(name).assign(xstar(name) + field_out - self.predictor_field_in)
         else:
             # Standard transport
             scheme.apply(xp(name), xstar(name))

From f6efea2cccb1368de90386a7f1bf8ba59ce09837 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Sat, 26 Oct 2024 13:17:19 +0100
Subject: [PATCH 9/9] revert example

---
 examples/shallow_water/williamson_5.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py
index 575eeef08..10fa2414a 100644
--- a/examples/shallow_water/williamson_5.py
+++ b/examples/shallow_water/williamson_5.py
@@ -84,15 +84,12 @@ def williamson_5(
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
-    transported_fields = [
-        TrapeziumRule(domain, "u"),
-        SSPRK3(domain, "D", subcycle_by_courant=0.3)
-    ]
+    transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")]
     transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods, predictor='D'
+        eqns, io, transported_fields, transport_methods
     )
 
     # ------------------------------------------------------------------------ #