From 2a929f8a62443b283c4622504f6e6cd602252cd7 Mon Sep 17 00:00:00 2001
From: Alex Brown <81297297+atb1995@users.noreply.github.com>
Date: Tue, 29 Oct 2024 13:40:48 +0000
Subject: [PATCH] Time Discretisation Improvements (#560)

---
 gusto/core/labels.py                          |   1 +
 gusto/equations/prognostic_equations.py       |   7 +-
 .../explicit_runge_kutta.py                   |  19 ++--
 gusto/time_discretisation/imex_runge_kutta.py | 107 +++++++++++++-----
 .../implicit_runge_kutta.py                   |  28 ++---
 .../multi_level_schemes.py                    |   8 ++
 gusto/time_discretisation/sdc.py              |   4 +
 .../time_discretisation.py                    |  42 +++++--
 8 files changed, 153 insertions(+), 63 deletions(-)

diff --git a/gusto/core/labels.py b/gusto/core/labels.py
index 7936926e2..15b687f65 100644
--- a/gusto/core/labels.py
+++ b/gusto/core/labels.py
@@ -102,6 +102,7 @@ def __call__(self, target, value=None):
 
 # labels for terms in the equations
 time_derivative = Label("time_derivative")
+nonlinear_time_derivative = Label("nonlinear_time_derivative")
 transport = Label("transport",
                   validator=lambda value: type(value) == TransportEquationType)
 diffusion = Label("diffusion")
diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py
index bf8dd7845..c8c708b3c 100644
--- a/gusto/equations/prognostic_equations.py
+++ b/gusto/equations/prognostic_equations.py
@@ -10,7 +10,8 @@
     replace_subject, replace_trial_function
 )
 from gusto.core import PrescribedFields
-from gusto.core.labels import time_derivative, prognostic, linearisation, mass_weighted
+from gusto.core.labels import (nonlinear_time_derivative, time_derivative,
+                               prognostic, linearisation, mass_weighted)
 from gusto.equations.common_forms import (
     advection_form, continuity_form, tracer_conservative_form
 )
@@ -163,8 +164,8 @@ def generate_mass_terms(self):
                         ref_density_idx = self.field_names.index(self.active_tracers[j].density_name)
                         ref_density = split(self.X)[ref_density_idx]
                         q = prog*ref_density
-                        mass_weighted_form = time_derivative(subject(prognostic(inner(q, test)*dx,
-                                                             field_name), self.X))
+                        mass_weighted_form = nonlinear_time_derivative(time_derivative(
+                            subject(prognostic(inner(q, test)*dx, field_name), self.X)))
 
                         mass = mass_weighted(standard_mass_form, mass_weighted_form)
             if i == 0:
diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py
index 88441eed6..dabdf93ec 100644
--- a/gusto/time_discretisation/explicit_runge_kutta.py
+++ b/gusto/time_discretisation/explicit_runge_kutta.py
@@ -160,9 +160,6 @@ def solver(self):
             return super().solver
 
         elif self.rk_formulation == RungeKuttaFormulation.predictor:
-            # In this case, don't set snes_type to ksp only, as we do want the
-            # outer Newton iteration. This is achieved by not calling the
-            # "super" method, in which the default snes_type is set to ksp_only
             solver_list = []
 
             for stage in range(self.nStages):
@@ -180,9 +177,6 @@ def solver(self):
             return solver_list
 
         elif self.rk_formulation == RungeKuttaFormulation.linear:
-            # In this case, don't set snes_type to ksp only, as we do want the
-            # outer Newton iteration. This is achieved by not calling the
-            # "super" method, in which the default snes_type is set to ksp_only
             problem = NonlinearVariationalProblem(
                 self.lhs - self.rhs[0], self.x1, bcs=self.bcs
             )
@@ -358,6 +352,10 @@ def solve_stage(self, x0, stage):
                 evaluate(self.x1, self.dt)
             if self.limiter is not None:
                 self.limiter.apply(self.x1)
+
+            # Set initial guess for solver
+            if stage > 0:
+                self.x_out.assign(self.k[stage-1])
             self.solver.solve()
 
             self.k[stage].assign(self.x_out)
@@ -376,8 +374,8 @@ def solve_stage(self, x0, stage):
             if stage == 0:
                 self.field_i[0].assign(x0)
 
-            # Use x0 as a first guess (otherwise may not converge)
-            self.field_i[stage+1].assign(x0)
+            # Use previous stage value as a first guess (otherwise may not converge)
+            self.field_i[stage+1].assign(self.field_i[stage])
 
             # Update field_i for physics / limiters
             for evaluate in self.evaluate_source:
@@ -423,6 +421,8 @@ def solve_stage(self, x0, stage):
                 if self.limiter is not None:
                     self.limiter.apply(self.field_rhs)
 
+                # Use previous stage value as a first guess (otherwise may not converge)
+                self.x1.assign(self.field_lhs[cycle_stage])
                 # Solve problem, placing solution in self.x1
                 self.solver[0].solve()
 
@@ -445,7 +445,8 @@ def solve_stage(self, x0, stage):
                     evaluate(self.field_rhs, self.original_dt)
                 if self.limiter is not None:
                     self.limiter.apply(self.field_rhs)
-
+                # Use x0 as a first guess (otherwise may not converge)
+                self.x1.assign(x0)
                 # Solve problem, placing solution in self.x1
                 self.solver[1].solve()
 
diff --git a/gusto/time_discretisation/imex_runge_kutta.py b/gusto/time_discretisation/imex_runge_kutta.py
index 0163d4121..75f27ae04 100644
--- a/gusto/time_discretisation/imex_runge_kutta.py
+++ b/gusto/time_discretisation/imex_runge_kutta.py
@@ -60,7 +60,8 @@ class IMEXRungeKutta(TimeDiscretisation):
     # --------------------------------------------------------------------------
 
     def __init__(self, domain, butcher_imp, butcher_exp, field_name=None,
-                 solver_parameters=None, limiter=None, options=None):
+                 linear_solver_parameters=None, nonlinear_solver_parameters=None,
+                 limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -73,20 +74,39 @@ def __init__(self, domain, butcher_imp, butcher_exp, field_name=None,
                 Runge Kutta time discretisation.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
+            linear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying linear solver. Defaults to None.
+            nonlinear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying nonlinear solver. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
         """
         super().__init__(domain, field_name=field_name,
-                         solver_parameters=solver_parameters,
+                         solver_parameters=nonlinear_solver_parameters,
                          options=options)
         self.butcher_imp = butcher_imp
         self.butcher_exp = butcher_exp
         self.nStages = int(np.shape(self.butcher_imp)[1])
 
+        # Set default linear and nonlinear solver options if none passed in
+        if linear_solver_parameters is None:
+            self.linear_solver_parameters = {'snes_type': 'ksponly',
+                                             'ksp_type': 'cg',
+                                             'pc_type': 'bjacobi',
+                                             'sub_pc_type': 'ilu'}
+        else:
+            self.linear_solver_parameters = linear_solver_parameters
+
+        if nonlinear_solver_parameters is None:
+            self.nonlinear_solver_parameters = {'snes_type': 'newtonls',
+                                                'ksp_type': 'gmres',
+                                                'pc_type': 'bjacobi',
+                                                'sub_pc_type': 'ilu'}
+        else:
+            self.nonlinear_solver_parameters = nonlinear_solver_parameters
+
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
         Set up the time discretisation based on the equation.
@@ -200,7 +220,7 @@ def solvers(self):
             # setup solver using residual defined in derived class
             problem = NonlinearVariationalProblem(self.res(stage), self.x_out, bcs=self.bcs)
             solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage)
-            solvers.append(NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name))
+            solvers.append(NonlinearVariationalSolver(problem, solver_parameters=self.nonlinear_solver_parameters, options_prefix=solver_name))
         return solvers
 
     @cached_property
@@ -209,19 +229,31 @@ def final_solver(self):
         # setup solver using lhs and rhs defined in derived class
         problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs)
         solver_name = self.field_name+self.__class__.__name__
-        return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
+        return NonlinearVariationalSolver(problem, solver_parameters=self.linear_solver_parameters, options_prefix=solver_name)
 
     @wrapper_apply
     def apply(self, x_out, x_in):
         self.x1.assign(x_in)
+        self.x_out.assign(x_in)
         solver_list = self.solvers
 
         for stage in range(self.nStages):
             self.solver = solver_list[stage]
+            # Set initial solver guess
+            if (stage > 0):
+                self.x_out.assign(self.xs[stage-1])
             self.solver.solve()
+
+            # Apply limiter
+            if self.limiter is not None:
+                self.limiter.apply(self.x_out)
             self.xs[stage].assign(self.x_out)
 
         self.final_solver.solve()
+
+        # Apply limiter
+        if self.limiter is not None:
+            self.limiter.apply(self.x_out)
         x_out.assign(self.x_out)
 
 
@@ -236,7 +268,8 @@ class IMEX_Euler(IMEXRungeKutta):
     y_1 = y^n + dt*F[y_1] + dt*S[y_0]                                         \n
     y^(n+1) = y^n + dt*F[y_1] + dt*S[y_0]
     """
-    def __init__(self, domain, field_name=None, solver_parameters=None,
+    def __init__(self, domain, field_name=None,
+                 linear_solver_parameters=None, nonlinear_solver_parameters=None,
                  limiter=None, options=None):
         """
         Args:
@@ -244,8 +277,10 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
+            linear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying linear solver. Defaults to None.
+            nonlinear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying nonlinear solver. Defaults to None.
             limiter (:class:`Limiter` object, optional): a limiter to apply to
                 the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
@@ -256,7 +291,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]])
         butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
-                         solver_parameters=solver_parameters,
+                         linear_solver_parameters=linear_solver_parameters,
+                         nonlinear_solver_parameters=nonlinear_solver_parameters,
                          limiter=limiter, options=options)
 
 
@@ -276,7 +312,8 @@ class IMEX_ARS3(IMEXRungeKutta):
     y^(n+1) = y^n + dt*(g*F[y_1]+(1-g)*F[y_2])                                \n
                   + dt*(0.5*S[y_1]+0.5*S[y_2])
     """
-    def __init__(self, domain, field_name=None, solver_parameters=None,
+    def __init__(self, domain, field_name=None,
+                 linear_solver_parameters=None, nonlinear_solver_parameters=None,
                  limiter=None, options=None):
         """
         Args:
@@ -284,8 +321,10 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
+            linear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying linear solver. Defaults to None.
+            nonlinear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying nonlinear solver. Defaults to None.
             limiter (:class:`Limiter` object, optional): a limiter to apply to
                 the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
@@ -298,7 +337,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]])
 
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
-                         solver_parameters=solver_parameters,
+                         linear_solver_parameters=linear_solver_parameters,
+                         nonlinear_solver_parameters=nonlinear_solver_parameters,
                          limiter=limiter, options=options)
 
 
@@ -318,15 +358,19 @@ class IMEX_ARK2(IMEXRungeKutta):
     y^(n+1) = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2])                           \n
                   + dt*(d*S[y_0]+d*S[y_1]+g*S[y_2])
     """
-    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None):
+    def __init__(self, domain, field_name=None,
+                 linear_solver_parameters=None, nonlinear_solver_parameters=None,
+                 limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
+            linear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying linear solver. Defaults to None.
+            nonlinear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying nonlinear solver. Defaults to None.
             limiter (:class:`Limiter` object, optional): a limiter to apply to
                 the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
@@ -340,7 +384,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None
         butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]])
         butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
-                         solver_parameters=solver_parameters,
+                         linear_solver_parameters=linear_solver_parameters,
+                         nonlinear_solver_parameters=nonlinear_solver_parameters,
                          limiter=limiter, options=options)
 
 
@@ -358,15 +403,19 @@ class IMEX_SSP3(IMEXRungeKutta):
     y^(n+1) = y^n + dt*(1/6*F[y_1]+1/6*F[y_2]+2/3*F[y_3])                     \n
                   + dt*(1/6*S[y_1]+1/6*S[y_2]+2/3*S[y_3])
     """
-    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None):
+    def __init__(self, domain, field_name=None,
+                 linear_solver_parameters=None, nonlinear_solver_parameters=None,
+                 limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
+            linear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying linear solver. Defaults to None.
+            nonlinear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying nonlinear solver. Defaults to None.
             limiter (:class:`Limiter` object, optional): a limiter to apply to
                 the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
@@ -378,7 +427,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None
         butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]])
         butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.), (1./6.), (2./3.)]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
-                         solver_parameters=solver_parameters,
+                         linear_solver_parameters=linear_solver_parameters,
+                         nonlinear_solver_parameters=nonlinear_solver_parameters,
                          limiter=limiter, options=options)
 
 
@@ -396,15 +446,19 @@ class IMEX_Trap2(IMEXRungeKutta):
     y_3 = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0]+0.5*S[y_2])        \n
     y^(n+1) = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0] + 0.5*S[y_2])  \n
     """
-    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None):
+    def __init__(self, domain, field_name=None,
+                 linear_solver_parameters=None, nonlinear_solver_parameters=None,
+                 limiter=None, options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            solver_parameters (dict, optional): dictionary of parameters to
-                pass to the underlying solver. Defaults to None.
+            linear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying linear solver. Defaults to None.
+            nonlinear_solver_parameters (dict, optional): dictionary of parameters to
+                pass to the underlying nonlinear solver. Defaults to None.
             limiter (:class:`Limiter` object, optional): a limiter to apply to
                 the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
@@ -416,5 +470,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None
         butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]])
         butcher_exp = np.array([[0., 0., 0., 0.], [1., 0., 0., 0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
-                         solver_parameters=solver_parameters,
+                         linear_solver_parameters=linear_solver_parameters,
+                         nonlinear_solver_parameters=nonlinear_solver_parameters,
                          limiter=limiter, options=options)
diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py
index 18b2368bb..8eb1c75ab 100644
--- a/gusto/time_discretisation/implicit_runge_kutta.py
+++ b/gusto/time_discretisation/implicit_runge_kutta.py
@@ -56,7 +56,7 @@ class ImplicitRungeKutta(TimeDiscretisation):
     # ---------------------------------------------------------------------------
 
     def __init__(self, domain, butcher_matrix, field_name=None,
-                 solver_parameters=None, limiter=None, options=None,):
+                 solver_parameters=None, options=None,):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -68,8 +68,6 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                 Defaults to None.
             solver_parameters (dict, optional): dictionary of parameters to
                 pass to the underlying solver. Defaults to None.
-            limiter (:class:`Limiter` object, optional): a limiter to apply to
-                the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
@@ -77,7 +75,7 @@ def __init__(self, domain, butcher_matrix, field_name=None,
         """
         super().__init__(domain, field_name=field_name,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options)
+                         options=options)
         self.butcher_matrix = butcher_matrix
         self.nStages = int(np.shape(self.butcher_matrix)[1])
 
@@ -131,15 +129,16 @@ def solve_stage(self, x0, stage):
         for i in range(stage):
             self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i])
 
-        if self.limiter is not None:
-            self.limiter.apply(self.x1)
-
         if self.idx is None and len(self.fs) > 1:
             self.xnph = tuple([self.dt*self.butcher_matrix[stage, stage]*a + b
                                for a, b in zip(split(self.x_out), split(self.x1))])
         else:
             self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out
         solver = self.solvers[stage]
+        # Set initial guess for solver
+        if (stage > 0):
+            self.x_out.assign(self.k[stage-1])
+
         solver.solve()
 
         self.k[stage].assign(self.x_out)
@@ -154,9 +153,6 @@ def apply(self, x_out, x_in):
         for i in range(self.nStages):
             x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i])
 
-        if self.limiter is not None:
-            self.limiter.apply(x_out)
-
 
 class ImplicitMidpoint(ImplicitRungeKutta):
     u"""
@@ -169,7 +165,7 @@ class ImplicitMidpoint(ImplicitRungeKutta):
     y^(n+1) = y^n + dt*k0                                                     \n
     """
     def __init__(self, domain, field_name=None, solver_parameters=None,
-                 limiter=None, options=None):
+                 options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -178,8 +174,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 Defaults to None.
             solver_parameters (dict, optional): dictionary of parameters to
                 pass to the underlying solver. Defaults to None.
-            limiter (:class:`Limiter` object, optional): a limiter to apply to
-                the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
@@ -188,7 +182,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         butcher_matrix = np.array([[0.5], [1.]])
         super().__init__(domain, butcher_matrix, field_name,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options)
+                         options=options)
 
 
 class QinZhang(ImplicitRungeKutta):
@@ -203,7 +197,7 @@ class QinZhang(ImplicitRungeKutta):
     y^(n+1) = y^n + 0.5*dt*(k0 + k1)                                          \n
     """
     def __init__(self, domain, field_name=None, solver_parameters=None,
-                 limiter=None, options=None):
+                 options=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -212,8 +206,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 Defaults to None.
             solver_parameters (dict, optional): dictionary of parameters to
                 pass to the underlying solver. Defaults to None.
-            limiter (:class:`Limiter` object, optional): a limiter to apply to
-                the evolving field to enforce monotonicity. Defaults to None.
             options (:class:`AdvectionOptions`, optional): an object containing
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
@@ -222,4 +214,4 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]])
         super().__init__(domain, butcher_matrix, field_name,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options)
+                         options=options)
diff --git a/gusto/time_discretisation/multi_level_schemes.py b/gusto/time_discretisation/multi_level_schemes.py
index 47e3b9652..a11671cc4 100644
--- a/gusto/time_discretisation/multi_level_schemes.py
+++ b/gusto/time_discretisation/multi_level_schemes.py
@@ -144,6 +144,8 @@ def apply(self, x_out, *x_in):
 
         self.xnm1.assign(x_in[0])
         self.x1.assign(x_in[1])
+        # Set initial solver guess
+        self.x_out.assign(x_in[1])
         solver.solve()
         x_out.assign(self.x_out)
 
@@ -221,6 +223,8 @@ def apply(self, x_out, *x_in):
 
         self.xnm1.assign(x_in[0])
         self.x1.assign(x_in[1])
+        # Set initial solver guess
+        self.x_out.assign(x_in[1])
         solver.solve()
         x_out.assign(self.x_out)
 
@@ -352,6 +356,8 @@ def apply(self, x_out, *x_in):
 
         for n in range(self.nlevels):
             self.x[n].assign(x_in[n])
+        # Set initial solver guess
+        self.x_out.assign(x_in[-1])
         solver.solve()
         x_out.assign(self.x_out)
 
@@ -507,5 +513,7 @@ def apply(self, x_out, *x_in):
 
         for n in range(self.nlevels):
             self.x[n].assign(x_in[n])
+        # Set initial solver guess
+        self.x_out.assign(x_in[-1])
         solver.solve()
         x_out.assign(self.x_out)
diff --git a/gusto/time_discretisation/sdc.py b/gusto/time_discretisation/sdc.py
index 45ff55fcc..f77688331 100644
--- a/gusto/time_discretisation/sdc.py
+++ b/gusto/time_discretisation/sdc.py
@@ -526,7 +526,11 @@ def apply(self, x_out, x_in):
                     self.fUnodes[m-1].assign(self.Urhs)
                 self.compute_quad_final()
                 # Compute y_(n+1) = y_n + sum(j=1,M) q_j*F(y_j)
+                self.U_fin.assign(self.Unodes[-1])
                 self.solver_fin.solve()
+                # Apply limiter if required
+                if self.limiter is not None:
+                    self.limiter.apply(self.U_fin)
                 x_out.assign(self.U_fin)
             else:
                 # Take value at final quadrature node dtau_M
diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py
index fdfe53f1e..50243e310 100644
--- a/gusto/time_discretisation/time_discretisation.py
+++ b/gusto/time_discretisation/time_discretisation.py
@@ -17,11 +17,11 @@
 from firedrake.utils import cached_property
 
 from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions
-from gusto.core.labels import time_derivative, prognostic, physics_label, mass_weighted
+from gusto.core.labels import (time_derivative, prognostic, physics_label,
+                               mass_weighted, nonlinear_time_derivative)
 from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.time_discretisation.wrappers import *
 
-
 __all__ = ["TimeDiscretisation", "ExplicitTimeDiscretisation", "BackwardEuler",
            "ThetaMethod", "TrapeziumRule", "TR_BDF2"]
 
@@ -110,7 +110,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
 
         # get default solver options if none passed in
         if solver_parameters is None:
-            self.solver_parameters = {'ksp_type': 'cg',
+            self.solver_parameters = {'ksp_type': 'gmres',
                                       'pc_type': 'bjacobi',
                                       'sub_pc_type': 'ilu'}
         else:
@@ -182,7 +182,6 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                             self.residual = self.residual.label_map(
                                 lambda t: t.get(prognostic) == field and t.has_label(mass_weighted),
                                 map_if_true=lambda t: t.get(mass_weighted))
-
         # -------------------------------------------------------------------- #
         # Set up Wrappers
         # -------------------------------------------------------------------- #
@@ -351,6 +350,15 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None,
         self.fixed_subcycles = fixed_subcycles
         self.subcycle_by_courant = subcycle_by_courant
 
+        # get default solver options if none passed in
+        if solver_parameters is None:
+            self.solver_parameters = {'snes_type': 'ksponly',
+                                      'ksp_type': 'cg',
+                                      'pc_type': 'bjacobi',
+                                      'sub_pc_type': 'ilu'}
+        else:
+            self.solver_parameters = solver_parameters
+
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
         Set up the time discretisation based on the equation.
@@ -375,6 +383,19 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         self.x0 = Function(self.fs)
         self.x1 = Function(self.fs)
 
+        # If the time_derivative term is nonlinear, we must use a nonlinear solver
+        if (
+            len(self.residual.label_map(
+                lambda t: t.has_label(nonlinear_time_derivative),
+                map_if_false=drop
+            )) > 0 and self.solver_parameters.get('snes_type') == 'ksponly'
+        ):
+            message = ('Switching to newton line search'
+                       + f' nonlinear solver for {self.field_name}'
+                       + ' as the time derivative term is nonlinear')
+            logger.warning(message)
+            self.solver_parameters['snes_type'] = 'newtonls'
+
     @cached_property
     def lhs(self):
         """Set up the discretisation's left hand side (the time derivative)."""
@@ -391,8 +412,6 @@ def solver(self):
         # setup linear solver using lhs and rhs defined in derived class
         problem = NonlinearVariationalProblem(self.lhs - self.rhs, self.x_out, bcs=self.bcs)
         solver_name = self.field_name+self.__class__.__name__
-        # If snes_type not specified by user, set this to ksp only to avoid outer Newton iteration
-        self.solver_parameters.setdefault('snes_type', 'ksponly')
         return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters,
                                           options_prefix=solver_name)
 
@@ -501,6 +520,8 @@ def apply(self, x_out, x_in):
             self.x_out.assign(x_in)
 
         self.x1.assign(x_in)
+        # Set initial solver guess
+        self.x_out.assign(x_in)
         self.solver.solve()
         x_out.assign(self.x_out)
 
@@ -585,6 +606,8 @@ def apply(self, x_out, x_in):
             x_in (:class:`Function`): the input field.
         """
         self.x1.assign(x_in)
+        # Set initial solver guess
+        self.x_out.assign(x_in)
         self.solver.solve()
         x_out.assign(self.x_out)
 
@@ -619,7 +642,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                          options=options)
 
 
-# TODO: this should be implemented as an ImplicitRK
 class TR_BDF2(TimeDiscretisation):
     """
     Implements the two stage implicit TR-BDF2 time stepping method, with a
@@ -747,6 +769,12 @@ def apply(self, x_out, x_in):
             x_in (:class:`Function`): the input field(s).
         """
         self.xn.assign(x_in)
+
+        # Set initial solver guess
+        self.xnpg.assign(x_in)
         self.solver_tr.solve()
+
+        # Set initial solver guess
+        self.x_out.assign(self.xnpg)
         self.solver_bdf2.solve()
         x_out.assign(self.x_out)