From c5d63857827c89177cc476149541babde9eb5463 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Wed, 14 Aug 2024 10:08:34 +0100
Subject: [PATCH 01/10] Accelerator and relaxation parameters added to improve
 solver convergence

---
 gusto/solvers/linear_solvers.py               | 110 +++++++++++-------
 .../semi_implicit_quasi_newton.py             |  31 ++++-
 2 files changed, 100 insertions(+), 41 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 4cf746cb4..b0823eafe 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -33,13 +33,19 @@
 class TimesteppingSolver(object, metaclass=ABCMeta):
     """Base class for timestepping linear solvers for Gusto."""
 
-    def __init__(self, equations, alpha=0.5, solver_parameters=None,
-                 overwrite_solver_parameters=False):
+    def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
+                 solver_parameters=None, overwrite_solver_parameters=False):
         """
         Args:
             equations (:class:`PrognosticEquation`): the model's equation.
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
+            tau_u (float, optional): the semi-implicit relaxation parameter for
+                wind. Defaults to 0.5.
+            tau_r (float, optional): the semi-implicit relaxation parameter for
+                density-like variable. Defaults to 0.5.
+            tau_t (float, optional): the semi-implicit relaxation parameter for
+                temperature-like variable. Defaults to 0.5.
             solver_parameters (dict, optional): contains the options to be
                 passed to the underlying :class:`LinearVariationalSolver`.
                 Defaults to None.
@@ -51,6 +57,10 @@ def __init__(self, equations, alpha=0.5, solver_parameters=None,
         self.equations = equations
         self.dt = equations.domain.dt
         self.alpha = alpha
+        # Set relaxation parameters.
+        self.tau_u=tau_u
+        self.tau_t=tau_t
+        self.tau_r=tau_r
 
         if solver_parameters is not None:
             if not overwrite_solver_parameters:
@@ -135,7 +145,7 @@ class CompressibleSolver(TimesteppingSolver):
                                                            'pc_type': 'bjacobi',
                                                            'sub_pc_type': 'ilu'}}}
 
-    def __init__(self, equations, alpha=0.5,
+    def __init__(self, equations, alpha=0.5, tau_u = 0.5, tau_r = 0.5, tau_t = 0.5,
                  quadrature_degree=None, solver_parameters=None,
                  overwrite_solver_parameters=False):
         """
@@ -143,6 +153,12 @@ def __init__(self, equations, alpha=0.5,
             equations (:class:`PrognosticEquation`): the model's equation.
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
+            tau_u (float, optional): the semi-implicit relaxation parameter for
+                wind. Defaults to 0.5.
+            tau_r (float, optional): the semi-implicit relaxation parameter for
+                density-like variable. Defaults to 0.5.
+            tau_t (float, optional): the semi-implicit relaxation parameter for
+                temperature-like variable. Defaults to 0.5.
             quadrature_degree (tuple, optional): a tuple (q_h, q_v) where q_h is
                 the required quadrature degree in the horizontal direction and
                 q_v is that in the vertical direction. Defaults to None.
@@ -164,7 +180,7 @@ def __init__(self, equations, alpha=0.5,
                 logger.warning("default quadrature degree most likely not sufficient for this degree element")
             self.quadrature_degree = (5, 5)
 
-        super().__init__(equations, alpha, solver_parameters,
+        super().__init__(equations, alpha, tau_u, tau_r, tau_t, solver_parameters,
                          overwrite_solver_parameters)
 
     @timed_function("Gusto:SolverSetup")
@@ -172,7 +188,9 @@ def _setup_solver(self):
 
         equations = self.equations
         dt = self.dt
-        beta_ = dt*self.alpha
+        beta_u_ = dt*self.tau_u
+        beta_t_ = dt*self.tau_t
+        beta_r_ = dt*self.tau_r
         cp = equations.parameters.cp
         Vu = equations.domain.spaces("HDiv")
         Vu_broken = FunctionSpace(equations.domain.mesh, BrokenElement(Vu.ufl_element()))
@@ -180,8 +198,10 @@ def _setup_solver(self):
         Vrho = equations.domain.spaces("DG")
 
         # Store time-stepping coefficients as UFL Constants
-        beta = Constant(beta_)
-        beta_cp = Constant(beta_ * cp)
+        beta_u = Constant(beta_u_)
+        beta_t = Constant(beta_t_)
+        beta_r = Constant(beta_r_)
+        beta_u_cp = Constant(beta_u * cp)
 
         h_deg = Vrho.ufl_element().degree()[0]
         v_deg = Vrho.ufl_element().degree()[1]
@@ -206,7 +226,7 @@ def _setup_solver(self):
 
         # Analytical (approximate) elimination of theta
         k = equations.domain.k             # Upward pointing unit vector
-        theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in
+        theta = -dot(k, u)*dot(k, grad(thetabar))*beta_t + theta_in
 
         # Only include theta' (rather than exner') in the vertical
         # component of the gradient
@@ -285,21 +305,21 @@ def L_tr(f):
         eqn = (
             # momentum equation
             u_mass
-            - beta_cp*div(theta_w*V(w))*exnerbar*dxp
+            - beta_u_cp*div(theta_w*V(w))*exnerbar*dxp
             # following does nothing but is preserved in the comments
             # to remind us why (because V(w) is purely vertical).
             # + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp
-            + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_hp
-            + beta_cp*dot(theta_w*V(w), n)*exnerbar_avg*ds_tbp
-            - beta_cp*div(thetabar_w*w)*exner*dxp
+            + beta_u_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_hp
+            + beta_u_cp*dot(theta_w*V(w), n)*exnerbar_avg*ds_tbp
+            - beta_u_cp*div(thetabar_w*w)*exner*dxp
             # trace terms appearing after integrating momentum equation
-            + beta_cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_vp + dS_hp)
-            + beta_cp*dot(thetabar_w*w, n)*l0*(ds_tbp + ds_vp)
+            + beta_u_cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_vp + dS_hp)
+            + beta_u_cp*dot(thetabar_w*w, n)*l0*(ds_tbp + ds_vp)
             # mass continuity equation
-            + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
-            + beta*jump(phi*u, n=n)*rhobar_avg('+')*(dS_v + dS_h)
+            + (phi*(rho - rho_in) - beta_r*inner(grad(phi), u)*rhobar)*dx
+            + beta_r*jump(phi*u, n=n)*rhobar_avg('+')*(dS_v + dS_h)
             # term added because u.n=0 is enforced weakly via the traces
-            + beta*phi*dot(u, n)*rhobar_avg*(ds_tb + ds_v)
+            + beta_r*phi*dot(u, n)*rhobar_avg*(ds_tb + ds_v)
             # constraint equation to enforce continuity of the velocity
             # through the interior facets and weakly impose the no-slip
             # condition
@@ -342,7 +362,7 @@ def L_tr(f):
 
         self.theta = Function(Vtheta)
         theta_eqn = gamma*(theta - theta_in
-                           + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta)*dx
+                           + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta_t)*dx
 
         theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta)
         self.theta_solver = LinearVariationalSolver(theta_problem,
@@ -446,13 +466,17 @@ def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
 
         dt = self.dt
-        beta_ = dt*self.alpha
+        beta_u_ = dt*self.tau_u
+        beta_p_ = dt*self.tau_r
+        beta_b_ = dt*self.tau_t
         Vu = equation.domain.spaces("HDiv")
         Vb = equation.domain.spaces("theta")
         Vp = equation.domain.spaces("DG")
 
         # Store time-stepping coefficients as UFL Constants
-        beta = Constant(beta_)
+        beta_u = Constant(beta_u_)
+        beta_p = Constant(beta_p_)
+        beta_b = Constant(beta_b_)
 
         # Split up the rhs vector (symbolically)
         self.xrhs = Function(self.equations.function_space)
@@ -468,7 +492,7 @@ def _setup_solver(self):
 
         # Analytical (approximate) elimination of theta
         k = equation.domain.k             # Upward pointing unit vector
-        b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in
+        b = -dot(k, u)*dot(k, grad(bbar))*beta_b + b_in
 
         # vertical projection
         def V(u):
@@ -476,13 +500,13 @@ def V(u):
 
         eqn = (
             inner(w, (u - u_in))*dx
-            - beta*div(w)*p*dx
-            - beta*inner(w, k)*b*dx
+            - beta_u*div(w)*p*dx
+            - beta_u*inner(w, k)*b*dx
         )
 
         if equation.compressible:
             cs = equation.parameters.cs
-            eqn += phi * (p - p_in) * dx + beta * phi * cs**2 * div(u) * dx
+            eqn += phi * (p - p_in) * dx + beta_p * phi * cs**2 * div(u) * dx
         else:
             eqn += phi * div(u) * dx
 
@@ -519,7 +543,7 @@ def trace_nullsp(T):
         self.b = Function(Vb)
 
         b_eqn = gamma*(b - b_in
-                       + dot(k, u)*dot(k, grad(bbar))*beta)*dx
+                       + dot(k, u)*dot(k, grad(bbar))*beta_b)*dx
 
         b_problem = LinearVariationalProblem(lhs(b_eqn),
                                              rhs(b_eqn),
@@ -589,7 +613,9 @@ class ThermalSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_ = dt*self.alpha
+        beta_u_ = dt*self.tau_u
+        beta_d_ = dt*self.tau_r
+        beta_b_ = dt*self.tau_t
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
         Vb = equation.domain.spaces("DG")
@@ -599,7 +625,9 @@ def _setup_solver(self):
             raise NotImplementedError("Field 'b' must exist to use the thermal linear solver in the SIQN scheme")
 
         # Store time-stepping coefficients as UFL Constants
-        beta = Constant(beta_)
+        beta_u = Constant(beta_u_)
+        beta_d = Constant(beta_d_)
+        beta_b = Constant(beta_b_)
 
         # Split up the rhs vector
         self.xrhs = Function(self.equations.function_space)
@@ -617,20 +645,20 @@ def _setup_solver(self):
         bbar = split(equation.X_ref)[2]
 
         # Approximate elimination of b
-        b = -dot(u, grad(bbar))*beta + b_in
+        b = -dot(u, grad(bbar))*beta_b + b_in
 
         n = FacetNormal(equation.domain.mesh)
 
         eqn = (
             inner(w, (u - u_in)) * dx
-            - beta * (D - Dbar) * div(w*bbar) * dx
-            + beta * jump(w*bbar, n) * avg(D-Dbar) * dS
-            - beta * 0.5 * Dbar * bbar * div(w) * dx
-            - beta * 0.5 * Dbar * b * div(w) * dx
-            - beta * 0.5 * bbar * div(w*(D-Dbar)) * dx
-            + beta * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS
+            - beta_u * (D - Dbar) * div(w*bbar) * dx
+            + beta_u * jump(w*bbar, n) * avg(D-Dbar) * dS
+            - beta_u * 0.5 * Dbar * bbar * div(w) * dx
+            - beta_u * 0.5 * Dbar * b * div(w) * dx
+            - beta_u * 0.5 * bbar * div(w*(D-Dbar)) * dx
+            + beta_u * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS
             + inner(phi, (D - D_in)) * dx
-            + beta * phi * Dbar * div(u) * dx
+            + beta_d * phi * Dbar * div(u) * dx
         )
 
         aeqn = lhs(eqn)
@@ -660,7 +688,7 @@ def trace_nullsp(T):
         u, D = self.uD.subfunctions
         self.b = Function(Vb)
 
-        b_eqn = gamma*(b - b_in + inner(u, grad(bbar))*beta) * dx
+        b_eqn = gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx
 
         b_problem = LinearVariationalProblem(lhs(b_eqn),
                                              rhs(b_eqn),
@@ -814,12 +842,14 @@ class MoistConvectiveSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_ = dt*self.alpha
+        beta_u_ = dt*self.tau_u
+        beta_d_ = dt*self.tau_r
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
 
         # Store time-stepping coefficients as UFL Constants
-        beta = Constant(beta_)
+        beta_u = Constant(beta_u_)
+        beta_d = Constant(beta_d_)
 
         # Split up the rhs vector
         self.xrhs = Function(self.equations.function_space)
@@ -838,9 +868,9 @@ def _setup_solver(self):
 
         eqn = (
             inner(w, (u - u_in)) * dx
-            - beta * (D - Dbar) * div(w*g) * dx
+            - beta_u * (D - Dbar) * div(w*g) * dx
             + inner(phi, (D - D_in)) * dx
-            + beta * phi * Dbar * div(u) * dx
+            + beta_d * phi * Dbar * div(u) * dx
         )
 
         aeqn = lhs(eqn)
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index ec3c3d349..63d73fac6 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -35,7 +35,7 @@ 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):
+                 num_outer=2, num_inner=2, accelerator=False):
 
         """
         Args:
@@ -84,6 +84,8 @@ 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
         """
 
         self.num_outer = num_outer
@@ -120,10 +122,12 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
                      + f"physics scheme {parametrisation.label.label}")
 
         self.active_transport = []
+        self.transported_fields = []
         for scheme in transport_schemes:
             assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop"
             assert scheme.field_name in equation_set.field_names
             self.active_transport.append((scheme.field_name, scheme))
+            self.transported_fields.append(scheme.field_name)
             # Check that there is a corresponding transport method
             method_found = False
             for method in spatial_methods:
@@ -184,6 +188,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
 
     def _apply_bcs(self):
         """
@@ -302,6 +307,9 @@ def timestep(self):
                 with timed_stage("Apply forcing terms"):
                     logger.info(f'Semi-implicit Quasi Newton: Implicit forcing {(outer, inner)}')
                     self.forcing.apply(xp, xnp1, xrhs, "implicit")
+                    if (inner > 0 and self.accelerator):
+                        # Zero implicit forcing to accelerate solver convergence
+                        self.forcing.zero_forcing_terms(self.equation, xp, xrhs, self.transported_fields)
 
                 xrhs -= xnp1(self.field_name)
                 xrhs += xrhs_phys
@@ -477,3 +485,24 @@ def apply(self, x_in, x_nl, x_out, label):
 
         x_out.assign(x_in(self.field_name))
         x_out += self.xF
+    
+    def zero_forcing_terms(self, equation, x_in, x_out, transported_field_names):
+        """
+        Zero forcing term F(x) for non-wind transport.
+
+        This takes x_in and x_out, where                                      \n
+        x_out = x_in + scale*F(x_nl)                                          \n
+        for some field x_nl and sets x_out = x_in for all non-wind transport terms 
+
+        Args:
+            equation (:class:`PrognosticEquationSet`): the prognostic
+                equation set to be solved
+            x_in (:class:`FieldCreator`): the field to be incremented.
+            x_out (:class:`FieldCreator`): the output field to be updated.
+            transported_field_names (str): list of fields names for transported fields
+        """
+        for field_name in transported_field_names:
+            if field_name != 'u':
+                logger.info(f'SIQN: Zeroing Implicit forcing for {field_name}')
+                field_index = equation.field_names.index(field_name)
+                x_out.subfunctions[field_index].assign(x_in(field_name))

From bd76e75eb6f80b4c463f449b39cceb29dd9bedaa Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Wed, 14 Aug 2024 10:55:03 +0100
Subject: [PATCH 02/10] lint fixes

---
 gusto/solvers/linear_solvers.py                  | 9 +++++----
 gusto/timestepping/semi_implicit_quasi_newton.py | 4 ++--
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index b0823eafe..0608e5d27 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -57,10 +57,11 @@ def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
         self.equations = equations
         self.dt = equations.domain.dt
         self.alpha = alpha
+
         # Set relaxation parameters.
-        self.tau_u=tau_u
-        self.tau_t=tau_t
-        self.tau_r=tau_r
+        self.tau_u = tau_u
+        self.tau_t = tau_t
+        self.tau_r = tau_r
 
         if solver_parameters is not None:
             if not overwrite_solver_parameters:
@@ -145,7 +146,7 @@ class CompressibleSolver(TimesteppingSolver):
                                                            'pc_type': 'bjacobi',
                                                            'sub_pc_type': 'ilu'}}}
 
-    def __init__(self, equations, alpha=0.5, tau_u = 0.5, tau_r = 0.5, tau_t = 0.5,
+    def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
                  quadrature_degree=None, solver_parameters=None,
                  overwrite_solver_parameters=False):
         """
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index 63d73fac6..af2c8b1c6 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -485,14 +485,14 @@ def apply(self, x_in, x_nl, x_out, label):
 
         x_out.assign(x_in(self.field_name))
         x_out += self.xF
-    
+
     def zero_forcing_terms(self, equation, x_in, x_out, transported_field_names):
         """
         Zero forcing term F(x) for non-wind transport.
 
         This takes x_in and x_out, where                                      \n
         x_out = x_in + scale*F(x_nl)                                          \n
-        for some field x_nl and sets x_out = x_in for all non-wind transport terms 
+        for some field x_nl and sets x_out = x_in for all non-wind transport terms
 
         Args:
             equation (:class:`PrognosticEquationSet`): the prognostic

From b1205e9462737d2d1484d9717ca9b5bf0df9b5da Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Thu, 15 Aug 2024 09:42:50 +0100
Subject: [PATCH 03/10] Change in response to review

---
 gusto/solvers/linear_solvers.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 0608e5d27..7c59bc08c 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -33,7 +33,7 @@
 class TimesteppingSolver(object, metaclass=ABCMeta):
     """Base class for timestepping linear solvers for Gusto."""
 
-    def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
+    def __init__(self, equations, alpha=0.5, tau_u=None, tau_r=None, tau_t=None,
                  solver_parameters=None, overwrite_solver_parameters=False):
         """
         Args:
@@ -41,11 +41,11 @@ def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_u (float, optional): the semi-implicit relaxation parameter for
-                wind. Defaults to 0.5.
+                wind. Defaults to None.
             tau_r (float, optional): the semi-implicit relaxation parameter for
-                density-like variable. Defaults to 0.5.
+                density-like variable. Defaults to None.
             tau_t (float, optional): the semi-implicit relaxation parameter for
-                temperature-like variable. Defaults to 0.5.
+                temperature-like variable. Defaults to None.
             solver_parameters (dict, optional): contains the options to be
                 passed to the underlying :class:`LinearVariationalSolver`.
                 Defaults to None.
@@ -58,10 +58,22 @@ def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
         self.dt = equations.domain.dt
         self.alpha = alpha
 
-        # Set relaxation parameters.
-        self.tau_u = tau_u
-        self.tau_t = tau_t
-        self.tau_r = tau_r
+        # Set relaxation parameters. If an alternative has not been given, set
+        # to semi-implicit off-centering factor
+        if tau_u is not None:
+            self.tau_u = tau_u
+        else:
+            self.tau_u = alpha
+
+        if tau_t is not None:
+            self.tau_t = tau_t
+        else:
+            self.tau_t = alpha
+
+        if tau_r is not None:
+            self.tau_r = tau_r
+        else:
+            self.tau_r = alpha
 
         if solver_parameters is not None:
             if not overwrite_solver_parameters:

From 98a9537ee0eba8e0044cdba2676d4e9d81f9c88b Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Fri, 16 Aug 2024 10:13:05 +0100
Subject: [PATCH 04/10] Changes for review

---
 gusto/solvers/linear_solvers.py                 | 17 +++--------------
 .../timestepping/semi_implicit_quasi_newton.py  |  2 +-
 2 files changed, 4 insertions(+), 15 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 7c59bc08c..08d128045 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -60,20 +60,9 @@ def __init__(self, equations, alpha=0.5, tau_u=None, tau_r=None, tau_t=None,
 
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        if tau_u is not None:
-            self.tau_u = tau_u
-        else:
-            self.tau_u = alpha
-
-        if tau_t is not None:
-            self.tau_t = tau_t
-        else:
-            self.tau_t = alpha
-
-        if tau_r is not None:
-            self.tau_r = tau_r
-        else:
-            self.tau_r = alpha
+        self.tau_u = tau_u if tau_u is not None else alpha
+        self.tau_t = tau_t if tau_t is not None else alpha
+        self.tau_r = tau_r if tau_r is not None else alpha
 
         if solver_parameters is not None:
             if not overwrite_solver_parameters:
diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
index af2c8b1c6..1e524a100 100644
--- a/gusto/timestepping/semi_implicit_quasi_newton.py
+++ b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -503,6 +503,6 @@ def zero_forcing_terms(self, equation, x_in, x_out, transported_field_names):
         """
         for field_name in transported_field_names:
             if field_name != 'u':
-                logger.info(f'SIQN: Zeroing Implicit forcing for {field_name}')
+                logger.info(f'Semi-Implicit Quasi Newton: Zeroing implicit forcing for {field_name}')
                 field_index = equation.field_names.index(field_name)
                 x_out.subfunctions[field_index].assign(x_in(field_name))

From 1a66a65a1c96cfe27c254a12f06e44da3e21a20d Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Fri, 16 Aug 2024 11:18:59 +0100
Subject: [PATCH 05/10] Switch to dictionary tau with appropriate checks

---
 gusto/solvers/linear_solvers.py | 61 ++++++++++++++++-----------------
 1 file changed, 30 insertions(+), 31 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 08d128045..7e1de8fc7 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -33,19 +33,15 @@
 class TimesteppingSolver(object, metaclass=ABCMeta):
     """Base class for timestepping linear solvers for Gusto."""
 
-    def __init__(self, equations, alpha=0.5, tau_u=None, tau_r=None, tau_t=None,
+    def __init__(self, equations, alpha=0.5, tau_values=None,
                  solver_parameters=None, overwrite_solver_parameters=False):
         """
         Args:
             equations (:class:`PrognosticEquation`): the model's equation.
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
-            tau_u (float, optional): the semi-implicit relaxation parameter for
-                wind. Defaults to None.
-            tau_r (float, optional): the semi-implicit relaxation parameter for
-                density-like variable. Defaults to None.
-            tau_t (float, optional): the semi-implicit relaxation parameter for
-                temperature-like variable. Defaults to None.
+            tau_values (dict, optional): contains the semi-implicit relaxation
+                parameters. Defaults to None.
             solver_parameters (dict, optional): contains the options to be
                 passed to the underlying :class:`LinearVariationalSolver`.
                 Defaults to None.
@@ -57,12 +53,14 @@ def __init__(self, equations, alpha=0.5, tau_u=None, tau_r=None, tau_t=None,
         self.equations = equations
         self.dt = equations.domain.dt
         self.alpha = alpha
+        self.tau_values = tau_values
 
-        # Set relaxation parameters. If an alternative has not been given, set
-        # to semi-implicit off-centering factor
-        self.tau_u = tau_u if tau_u is not None else alpha
-        self.tau_t = tau_t if tau_t is not None else alpha
-        self.tau_r = tau_r if tau_r is not None else alpha
+        # Check that we have a tau value for each field
+        if self.tau_values is not None:
+            for field in equations.field_names:
+                if field not in self.tau_values:
+                    raise KeyError(f"You must select tau values for all fields to be solved, "
+                                   f"missing value is {field}")
 
         if solver_parameters is not None:
             if not overwrite_solver_parameters:
@@ -147,7 +145,7 @@ class CompressibleSolver(TimesteppingSolver):
                                                            'pc_type': 'bjacobi',
                                                            'sub_pc_type': 'ilu'}}}
 
-    def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
+    def __init__(self, equations, alpha=0.5, tau_values=None,
                  quadrature_degree=None, solver_parameters=None,
                  overwrite_solver_parameters=False):
         """
@@ -155,12 +153,8 @@ def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
             equations (:class:`PrognosticEquation`): the model's equation.
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
-            tau_u (float, optional): the semi-implicit relaxation parameter for
-                wind. Defaults to 0.5.
-            tau_r (float, optional): the semi-implicit relaxation parameter for
-                density-like variable. Defaults to 0.5.
-            tau_t (float, optional): the semi-implicit relaxation parameter for
-                temperature-like variable. Defaults to 0.5.
+            tau_values (dict, optional): contains the semi-implicit relaxation
+                parameters. Defaults to None.
             quadrature_degree (tuple, optional): a tuple (q_h, q_v) where q_h is
                 the required quadrature degree in the horizontal direction and
                 q_v is that in the vertical direction. Defaults to None.
@@ -182,7 +176,7 @@ def __init__(self, equations, alpha=0.5, tau_u=0.5, tau_r=0.5, tau_t=0.5,
                 logger.warning("default quadrature degree most likely not sufficient for this degree element")
             self.quadrature_degree = (5, 5)
 
-        super().__init__(equations, alpha, tau_u, tau_r, tau_t, solver_parameters,
+        super().__init__(equations, alpha, tau_values, solver_parameters,
                          overwrite_solver_parameters)
 
     @timed_function("Gusto:SolverSetup")
@@ -190,9 +184,12 @@ def _setup_solver(self):
 
         equations = self.equations
         dt = self.dt
-        beta_u_ = dt*self.tau_u
-        beta_t_ = dt*self.tau_t
-        beta_r_ = dt*self.tau_r
+        # Set relaxation parameters. If an alternative has not been given, set
+        # to semi-implicit off-centering factor
+        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
+        beta_t_ = dt*self.tau_values["theta"] if self.tau_values is not None else self.alpha
+        beta_r_ = dt*self.tau_values["rho"] if self.tau_values is not None else self.alpha
+
         cp = equations.parameters.cp
         Vu = equations.domain.spaces("HDiv")
         Vu_broken = FunctionSpace(equations.domain.mesh, BrokenElement(Vu.ufl_element()))
@@ -468,9 +465,11 @@ def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
 
         dt = self.dt
-        beta_u_ = dt*self.tau_u
-        beta_p_ = dt*self.tau_r
-        beta_b_ = dt*self.tau_t
+        # Set relaxation parameters. If an alternative has not been given, set
+        # to semi-implicit off-centering factor
+        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
+        beta_p_ = dt*self.tau_values["p"] if self.tau_values is not None else self.alpha
+        beta_b_ = dt*self.tau_values["b"] if self.tau_values is not None else self.alpha
         Vu = equation.domain.spaces("HDiv")
         Vb = equation.domain.spaces("theta")
         Vp = equation.domain.spaces("DG")
@@ -615,9 +614,9 @@ class ThermalSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_u
-        beta_d_ = dt*self.tau_r
-        beta_b_ = dt*self.tau_t
+        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
+        beta_d_ = dt*self.tau_values["D"] if self.tau_values is not None else self.alpha
+        beta_b_ = dt*self.tau_values["b"] if self.tau_values is not None else self.alpha
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
         Vb = equation.domain.spaces("DG")
@@ -844,8 +843,8 @@ class MoistConvectiveSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_u
-        beta_d_ = dt*self.tau_r
+        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
+        beta_d_ = dt*self.tau_values["D"] if self.tau_values is not None else self.alpha
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
 

From 6aef3210dee0b1e068e14ebd21531faf34317f88 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Fri, 16 Aug 2024 11:37:49 +0100
Subject: [PATCH 06/10] Fix to how the dictionary works

---
 gusto/solvers/linear_solvers.py | 37 +++++++++++++--------------------
 1 file changed, 15 insertions(+), 22 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 7e1de8fc7..4979f7662 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -33,7 +33,7 @@
 class TimesteppingSolver(object, metaclass=ABCMeta):
     """Base class for timestepping linear solvers for Gusto."""
 
-    def __init__(self, equations, alpha=0.5, tau_values=None,
+    def __init__(self, equations, alpha=0.5, tau_values={},
                  solver_parameters=None, overwrite_solver_parameters=False):
         """
         Args:
@@ -41,7 +41,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None,
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_values (dict, optional): contains the semi-implicit relaxation
-                parameters. Defaults to None.
+                parameters. Defaults to an empty dictionary.
             solver_parameters (dict, optional): contains the options to be
                 passed to the underlying :class:`LinearVariationalSolver`.
                 Defaults to None.
@@ -55,13 +55,6 @@ def __init__(self, equations, alpha=0.5, tau_values=None,
         self.alpha = alpha
         self.tau_values = tau_values
 
-        # Check that we have a tau value for each field
-        if self.tau_values is not None:
-            for field in equations.field_names:
-                if field not in self.tau_values:
-                    raise KeyError(f"You must select tau values for all fields to be solved, "
-                                   f"missing value is {field}")
-
         if solver_parameters is not None:
             if not overwrite_solver_parameters:
                 p = flatten_parameters(self.solver_parameters)
@@ -145,7 +138,7 @@ class CompressibleSolver(TimesteppingSolver):
                                                            'pc_type': 'bjacobi',
                                                            'sub_pc_type': 'ilu'}}}
 
-    def __init__(self, equations, alpha=0.5, tau_values=None,
+    def __init__(self, equations, alpha=0.5, tau_values={},
                  quadrature_degree=None, solver_parameters=None,
                  overwrite_solver_parameters=False):
         """
@@ -154,7 +147,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None,
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_values (dict, optional): contains the semi-implicit relaxation
-                parameters. Defaults to None.
+                parameters. Defaults to an empty dictionary.
             quadrature_degree (tuple, optional): a tuple (q_h, q_v) where q_h is
                 the required quadrature degree in the horizontal direction and
                 q_v is that in the vertical direction. Defaults to None.
@@ -186,9 +179,9 @@ def _setup_solver(self):
         dt = self.dt
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
-        beta_t_ = dt*self.tau_values["theta"] if self.tau_values is not None else self.alpha
-        beta_r_ = dt*self.tau_values["rho"] if self.tau_values is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
+        beta_t_ = dt*self.tau_values.get("theta") if self.tau_values.get("theta") is not None else self.alpha
+        beta_r_ = dt*self.tau_values.get("rho") if self.tau_values.get("rho") is not None else self.alpha
 
         cp = equations.parameters.cp
         Vu = equations.domain.spaces("HDiv")
@@ -467,9 +460,9 @@ def _setup_solver(self):
         dt = self.dt
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
-        beta_p_ = dt*self.tau_values["p"] if self.tau_values is not None else self.alpha
-        beta_b_ = dt*self.tau_values["b"] if self.tau_values is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
+        beta_p_ = dt*self.tau_values.get("p") if self.tau_values.get("p") is not None else self.alpha
+        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else self.alpha
         Vu = equation.domain.spaces("HDiv")
         Vb = equation.domain.spaces("theta")
         Vp = equation.domain.spaces("DG")
@@ -614,9 +607,9 @@ class ThermalSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
-        beta_d_ = dt*self.tau_values["D"] if self.tau_values is not None else self.alpha
-        beta_b_ = dt*self.tau_values["b"] if self.tau_values is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
+        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else self.alpha
+        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else self.alpha
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
         Vb = equation.domain.spaces("DG")
@@ -843,8 +836,8 @@ class MoistConvectiveSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_values["u"] if self.tau_values is not None else self.alpha
-        beta_d_ = dt*self.tau_values["D"] if self.tau_values is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
+        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else self.alpha
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
 

From c15f99b58e575bf1415274a40e4e239492598fd6 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Fri, 16 Aug 2024 11:40:54 +0100
Subject: [PATCH 07/10] Add in dt to fix

---
 gusto/solvers/linear_solvers.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 4979f7662..e3603bb2a 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -179,9 +179,9 @@ def _setup_solver(self):
         dt = self.dt
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
-        beta_t_ = dt*self.tau_values.get("theta") if self.tau_values.get("theta") is not None else self.alpha
-        beta_r_ = dt*self.tau_values.get("rho") if self.tau_values.get("rho") is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
+        beta_t_ = dt*self.tau_values.get("theta") if self.tau_values.get("theta") is not None else dt*self.alpha
+        beta_r_ = dt*self.tau_values.get("rho") if self.tau_values.get("rho") is not None else dt*self.alpha
 
         cp = equations.parameters.cp
         Vu = equations.domain.spaces("HDiv")
@@ -460,9 +460,9 @@ def _setup_solver(self):
         dt = self.dt
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
-        beta_p_ = dt*self.tau_values.get("p") if self.tau_values.get("p") is not None else self.alpha
-        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
+        beta_p_ = dt*self.tau_values.get("p") if self.tau_values.get("p") is not None else dt*self.alpha
+        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else dt*self.alpha
         Vu = equation.domain.spaces("HDiv")
         Vb = equation.domain.spaces("theta")
         Vp = equation.domain.spaces("DG")
@@ -607,9 +607,9 @@ class ThermalSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
-        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else self.alpha
-        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
+        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else dt*self.alpha
+        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else dt*self.alpha
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
         Vb = equation.domain.spaces("DG")
@@ -836,8 +836,8 @@ class MoistConvectiveSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else self.alpha
-        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else self.alpha
+        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
+        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else dt*self.alpha
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
 

From dc3840796a042af5ff8bf878d8132b1fc9132039 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown@metoffice.gov.uk>
Date: Fri, 16 Aug 2024 11:47:27 +0100
Subject: [PATCH 08/10] Remove empty dict option and tidy up dict get

---
 gusto/solvers/linear_solvers.py | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index e3603bb2a..07e39e26b 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -33,7 +33,7 @@
 class TimesteppingSolver(object, metaclass=ABCMeta):
     """Base class for timestepping linear solvers for Gusto."""
 
-    def __init__(self, equations, alpha=0.5, tau_values={},
+    def __init__(self, equations, alpha=0.5, tau_values=None,
                  solver_parameters=None, overwrite_solver_parameters=False):
         """
         Args:
@@ -41,7 +41,7 @@ def __init__(self, equations, alpha=0.5, tau_values={},
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_values (dict, optional): contains the semi-implicit relaxation
-                parameters. Defaults to an empty dictionary.
+                parameters. Defaults to None.
             solver_parameters (dict, optional): contains the options to be
                 passed to the underlying :class:`LinearVariationalSolver`.
                 Defaults to None.
@@ -53,7 +53,7 @@ def __init__(self, equations, alpha=0.5, tau_values={},
         self.equations = equations
         self.dt = equations.domain.dt
         self.alpha = alpha
-        self.tau_values = tau_values
+        self.tau_values = tau_values if tau_values is not None else {}
 
         if solver_parameters is not None:
             if not overwrite_solver_parameters:
@@ -138,7 +138,7 @@ class CompressibleSolver(TimesteppingSolver):
                                                            'pc_type': 'bjacobi',
                                                            'sub_pc_type': 'ilu'}}}
 
-    def __init__(self, equations, alpha=0.5, tau_values={},
+    def __init__(self, equations, alpha=0.5, tau_values=None,
                  quadrature_degree=None, solver_parameters=None,
                  overwrite_solver_parameters=False):
         """
@@ -147,7 +147,7 @@ def __init__(self, equations, alpha=0.5, tau_values={},
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_values (dict, optional): contains the semi-implicit relaxation
-                parameters. Defaults to an empty dictionary.
+                parameters. Defaults to None.
             quadrature_degree (tuple, optional): a tuple (q_h, q_v) where q_h is
                 the required quadrature degree in the horizontal direction and
                 q_v is that in the vertical direction. Defaults to None.
@@ -179,9 +179,9 @@ def _setup_solver(self):
         dt = self.dt
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
-        beta_t_ = dt*self.tau_values.get("theta") if self.tau_values.get("theta") is not None else dt*self.alpha
-        beta_r_ = dt*self.tau_values.get("rho") if self.tau_values.get("rho") is not None else dt*self.alpha
+        beta_u_ = dt*self.tau_values.get("u", self.alpha)
+        beta_t_ = dt*self.tau_values.get("theta", self.alpha)
+        beta_r_ = dt*self.tau_values.get("rho", self.alpha)
 
         cp = equations.parameters.cp
         Vu = equations.domain.spaces("HDiv")
@@ -460,9 +460,9 @@ def _setup_solver(self):
         dt = self.dt
         # Set relaxation parameters. If an alternative has not been given, set
         # to semi-implicit off-centering factor
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
-        beta_p_ = dt*self.tau_values.get("p") if self.tau_values.get("p") is not None else dt*self.alpha
-        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else dt*self.alpha
+        beta_u_ = dt*self.tau_values.get("u", self.alpha)
+        beta_p_ = dt*self.tau_values.get("p", self.alpha)
+        beta_b_ = dt*self.tau_values.get("b", self.alpha)
         Vu = equation.domain.spaces("HDiv")
         Vb = equation.domain.spaces("theta")
         Vp = equation.domain.spaces("DG")
@@ -607,9 +607,9 @@ class ThermalSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
-        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else dt*self.alpha
-        beta_b_ = dt*self.tau_values.get("b") if self.tau_values.get("b") is not None else dt*self.alpha
+        beta_u_ = dt*self.tau_values.get("u", self.alpha)
+        beta_d_ = dt*self.tau_values.get("D", self.alpha)
+        beta_b_ = dt*self.tau_values.get("b", self.alpha)
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
         Vb = equation.domain.spaces("DG")
@@ -836,8 +836,8 @@ class MoistConvectiveSWSolver(TimesteppingSolver):
     def _setup_solver(self):
         equation = self.equations      # just cutting down line length a bit
         dt = self.dt
-        beta_u_ = dt*self.tau_values.get("u") if self.tau_values.get("u") is not None else dt*self.alpha
-        beta_d_ = dt*self.tau_values.get("D") if self.tau_values.get("D") is not None else dt*self.alpha
+        beta_u_ = dt*self.tau_values.get("u", self.alpha)
+        beta_d_ = dt*self.tau_values.get("D", self.alpha)
         Vu = equation.domain.spaces("HDiv")
         VD = equation.domain.spaces("DG")
 

From 52074c07071cc8695477d1cdfef54727eb1c8556 Mon Sep 17 00:00:00 2001
From: Alex Brown <81297297+atb1995@users.noreply.github.com>
Date: Fri, 16 Aug 2024 12:06:14 +0100
Subject: [PATCH 09/10] Update gusto/solvers/linear_solvers.py

Co-authored-by: Dr Jemma Shipton <j.shipton@exeter.ac.uk>
---
 gusto/solvers/linear_solvers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 07e39e26b..34b699d27 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -41,7 +41,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None,
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_values (dict, optional): contains the semi-implicit relaxation
-                parameters. Defaults to None.
+                parameters. Defaults to None, in which case the value of alpha is used.
             solver_parameters (dict, optional): contains the options to be
                 passed to the underlying :class:`LinearVariationalSolver`.
                 Defaults to None.

From a88e3cb742eea53ca11d61855a643eb8dc34aaff Mon Sep 17 00:00:00 2001
From: Alex Brown <81297297+atb1995@users.noreply.github.com>
Date: Fri, 16 Aug 2024 12:06:20 +0100
Subject: [PATCH 10/10] Update gusto/solvers/linear_solvers.py

Co-authored-by: Dr Jemma Shipton <j.shipton@exeter.ac.uk>
---
 gusto/solvers/linear_solvers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py
index 34b699d27..f3a610da7 100644
--- a/gusto/solvers/linear_solvers.py
+++ b/gusto/solvers/linear_solvers.py
@@ -147,7 +147,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None,
             alpha (float, optional): the semi-implicit off-centring factor.
                 Defaults to 0.5. A value of 1 is fully-implicit.
             tau_values (dict, optional): contains the semi-implicit relaxation
-                parameters. Defaults to None.
+                parameters. Defaults to None, in which case the value of alpha is used.
             quadrature_degree (tuple, optional): a tuple (q_h, q_v) where q_h is
                 the required quadrature degree in the horizontal direction and
                 q_v is that in the vertical direction. Defaults to None.