From 5db19373c8660aa0d3a236baefb9f049ca11ad81 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Thu, 7 Sep 2023 10:38:05 +0100
Subject: [PATCH 01/13] Tidy of comments, ImplicitMidpoint and QinZhang added

---
 gusto/time_discretisation.py | 176 ++++++++++++++++++++++++++++++++---
 1 file changed, 165 insertions(+), 11 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 2310102a8..d14b43d02 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -21,9 +21,9 @@
 import numpy as np
 
 
-__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "SSPRK3", "RK4",
-           "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog",
-           "AdamsMoulton", "AdamsBashforth"]
+__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", 
+           "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", 
+           "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang"]
 
 
 def wrapper_apply(original_apply):
@@ -234,6 +234,121 @@ def apply(self, x_out, x_in):
         pass
 
 
+class ImplicitMultistage(TimeDiscretisation):
+    """
+    A class for implementing general diagonally implicit multistage (Runge-Kutta)
+    methods based on its Butcher tableau.
+
+    A Butcher tableau is formed in the following way for a s-th order 
+    diagonally implicit scheme:
+
+    c_0 | a_00 a_01  .    a_0s
+    c_1 | a_10 a_11       a_1s
+     .  |   .   .    .    .
+     .  |   .   .    .    .
+    c_s | a_s0 a_s1  .    a_ss
+     -------------------------
+        |  b_1  b_2  ...  b_s
+
+    The gradient function has no time-dependence, so c elements are not needed.
+    A square 'butcher_matrix' is defined in each class that uses
+    the ImplicitMultiStage structure,
+
+    [a_00   0   .       0  ]
+    [a_10 a_11  .       0  ]
+    [  .    .   .       .  ]
+    [ b_0  b_1  .       b_s]
+
+    Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods.
+
+    There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
+
+    For an s stage method,
+      At iteration i (from 0 to s-1)
+        An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i)
+        Compute the gradient at the intermediate location, k_i = F(y_i)
+
+    At the last stage, compute the new solution by:
+    y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i)
+
+    """
+
+    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None):
+        super().__init__(domain, field_name=field_name,
+                         solver_parameters=solver_parameters,
+                         limiter=limiter, options=options)
+        if butcher_matrix is not None:
+            self.butcher_matrix = butcher_matrix
+            self.nStages = int(np.shape(self.butcher_matrix)[1])
+
+    def setup(self, equation, apply_bcs=True, *active_labels):
+        """
+        Set up the time discretisation based on the equation.
+
+        Args:
+            equation (:class:`PrognosticEquation`): the model's equation.
+            *active_labels (:class:`Label`): labels indicating which terms of
+                the equation to include.
+        """
+
+        super().setup(equation, apply_bcs, *active_labels)
+
+        self.k = [Function(self.fs) for i in range(self.nStages)]
+
+    def lhs(self):
+        return super().lhs
+    
+    def rhs(self):
+        return super().rhs
+    
+    def solvers(self, stage):
+        residual = self.residual.label_map(
+            lambda t: t.has_label(time_derivative),
+            map_if_true=drop,
+            map_if_false=replace_subject(self.xnph, self.idx),
+        )
+        mass_form = self.residual.label_map(
+            lambda t: t.has_label(time_derivative),
+            map_if_false=drop)
+        residual += mass_form.label_map(all_terms,
+                                        replace_subject(self.x_out, self.idx))
+
+        problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs)
+        
+        solver_name = self.field_name+self.__class__.__name__+"%s"%(stage)
+        return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters,
+                                          options_prefix=solver_name)
+
+    def solve_stage(self, x0, stage):
+        self.x1.assign(x0)
+        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
+        
+        self.solvers(stage).solve()
+
+        self.k[stage].assign(self.x_out)
+
+    def apply(self, x_out, x_in):
+        
+        for i in range(self.nStages):
+            self.solve_stage(x_in, i)
+
+        x_out.assign(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 ExplicitTimeDiscretisation(TimeDiscretisation):
     """Base class for explicit time discretisations."""
 
@@ -342,23 +457,25 @@ class ExplicitMultistage(ExplicitTimeDiscretisation):
 
     A Butcher tableau is formed in the following way for a s-th order explicit scheme:
 
-    c_0 | a_00
-    c_1 | a_10 a_11
-     .  | a_20 a_21 a_22
+    c_0 | a_00 a_01  .    a_0s
+    c_1 | a_10 a_11       a_1s
+     .  |   .   .    .    .
      .  |   .   .    .    .
+    c_s | a_s0 a_s1  .    a_ss
      -------------------------
-    c_s |  b_1  b_2  ...  b_s
+        |  b_1  b_2  ...  b_s
 
     The gradient function has no time-dependence, so c elements are not needed.
     A square 'butcher_matrix' is defined in each class that uses
     the ExplicitMultiStage structure,
 
-    [a_00   0   .       0  ]
-    [a_10 a_11  .       0  ]
+    [a_10   0   .       0  ]
+    [a_20 a_21  .       0  ]
     [  .    .   .       .  ]
-    [ b_0  b_1  .   b_{s-1}]
+    [ b_0  b_1  .       b_s]
 
-    All upper diagonal a_ij elements are zero for explicit methods.
+    All upper diagonal a_ij elements are zero for explicit methods. We exclude the first
+    row of the butcher tableau from our butcher matrix as the row is always zeros.
 
     There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
 
@@ -1341,3 +1458,40 @@ def apply(self, x_out, *x_in):
             self.x[n].assign(x_in[n])
         solver.solve()
         x_out.assign(self.x_out)
+
+
+class ImplicitMidpoint(ImplicitMultistage):
+    u"""
+    Implements the Implicit Midpoint method as a 1-stage Runge Kutta method.
+
+    The method, for solving
+    ∂y/∂t = F(y), can be written as:
+
+    k0 = F[y^n + 0.5*dt*k0]
+    y^(n+1) = y^n + dt*k0
+    """
+    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None):
+        super().__init__(domain, field_name,
+                         solver_parameters=solver_parameters,
+                         limiter=limiter, options=options)
+        self.butcher_matrix = np.array([[0.5],[1.]])
+        self.nStages = int(np.shape(self.butcher_matrix)[1])
+
+
+class QinZhang(ImplicitMultistage):
+    u"""
+    Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method.
+
+    The method, for solving
+    ∂y/∂t = F(y), can be written as:
+
+    k0 = F[y^n + 0.25*dt*k0]
+    k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1]
+    y^(n+1) = y^n + 0.5*dt*(k0 + k1)
+    """
+    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None):
+        super().__init__(domain, field_name,
+                         solver_parameters=solver_parameters,
+                         limiter=limiter, options=options)
+        self.butcher_matrix = np.array([[0.25, 0],[0.5, 0.25],[0.5, 0.5]])
+        self.nStages = int(np.shape(self.butcher_matrix)[1])

From bb9155a7b580e7718ad0e07bf4e29fe1795dfd28 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Thu, 7 Sep 2023 10:50:42 +0100
Subject: [PATCH 02/13] Tidy of comments, ImplicitMidpoint and QinZhang added

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

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index d14b43d02..f25098193 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -7,7 +7,7 @@
 
 from abc import ABCMeta, abstractmethod, abstractproperty
 from firedrake import (Function, TestFunction, NonlinearVariationalProblem,
-                       NonlinearVariationalSolver, DirichletBC)
+                       NonlinearVariationalSolver, DirichletBC, split)
 from firedrake.formmanipulation import split_form
 from firedrake.utils import cached_property
 

From ec25b89502ba2a11f19a346ab061331568706b74 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Thu, 7 Sep 2023 11:11:29 +0100
Subject: [PATCH 03/13] Lint tidy

---
 gusto/time_discretisation.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index f25098193..dce51c963 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -21,8 +21,8 @@
 import numpy as np
 
 
-__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", 
-           "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", 
+__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage",
+           "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2",
            "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang"]
 
 
@@ -239,7 +239,7 @@ class ImplicitMultistage(TimeDiscretisation):
     A class for implementing general diagonally implicit multistage (Runge-Kutta)
     methods based on its Butcher tableau.
 
-    A Butcher tableau is formed in the following way for a s-th order 
+    A Butcher tableau is formed in the following way for a s-th order
     diagonally implicit scheme:
 
     c_0 | a_00 a_01  .    a_0s
@@ -297,10 +297,10 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
     def lhs(self):
         return super().lhs
-    
+
     def rhs(self):
         return super().rhs
-    
+
     def solvers(self, stage):
         residual = self.residual.label_map(
             lambda t: t.has_label(time_derivative),
@@ -314,37 +314,37 @@ def solvers(self, stage):
                                         replace_subject(self.x_out, self.idx))
 
         problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs)
-        
-        solver_name = self.field_name+self.__class__.__name__+"%s"%(stage)
+
+        solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage)
         return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters,
                                           options_prefix=solver_name)
 
     def solve_stage(self, x0, stage):
         self.x1.assign(x0)
         for i in range(stage):
-            self.x1.assign(self.x1 + self.butcher_matrix[stage,i]*self.dt*self.k[i])
-        
+            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))])
+            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
-        
+            self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out
+
         self.solvers(stage).solve()
 
         self.k[stage].assign(self.x_out)
 
     def apply(self, x_out, x_in):
-        
+
         for i in range(self.nStages):
             self.solve_stage(x_in, i)
 
         x_out.assign(x_in)
         for i in range(self.nStages):
-            x_out.assign(x_out + self.butcher_matrix[self.nStages,i]*self.dt*self.k[i])
-        
+            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)
 
@@ -1474,7 +1474,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None
         super().__init__(domain, field_name,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        self.butcher_matrix = np.array([[0.5],[1.]])
+        self.butcher_matrix = np.array([[0.5], [1.]])
         self.nStages = int(np.shape(self.butcher_matrix)[1])
 
 
@@ -1493,5 +1493,5 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None
         super().__init__(domain, field_name,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        self.butcher_matrix = np.array([[0.25, 0],[0.5, 0.25],[0.5, 0.5]])
+        self.butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]])
         self.nStages = int(np.shape(self.butcher_matrix)[1])

From aa9921a4ba7e071af91bd4cdcf76feb18cc8a3f7 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Thu, 7 Sep 2023 14:59:34 +0100
Subject: [PATCH 04/13] Testing added

---
 integration-tests/model/test_time_discretisation.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py
index 5be78c871..a61e43d9e 100644
--- a/integration-tests/model/test_time_discretisation.py
+++ b/integration-tests/model/test_time_discretisation.py
@@ -8,7 +8,7 @@ def run(timestepper, tmax, f_end):
     return norm(timestepper.fields("f") - f_end) / norm(f_end)
 
 
-@pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint",
+@pytest.mark.parametrize("scheme", ["ssprk", "TrapeziumRule", "ImplicitMidpoint", "QinZhang",
                                     "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", "AdamsMoulton"])
 def test_time_discretisation(tmpdir, scheme, tracer_setup):
     if (scheme == "AdamsBashforth"):
@@ -27,8 +27,12 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup):
 
     if scheme == "ssprk":
         transport_scheme = SSPRK3(domain)
-    elif scheme == "implicit_midpoint":
+    elif scheme == "TrapeziumRule":
         transport_scheme = TrapeziumRule(domain)
+    elif scheme == "ImplicitMidpoint":
+      transport_scheme = ImplicitMidpoint(domain)
+    elif scheme == "QinZhang":
+       transport_scheme = QinZhang(domain)
     elif scheme == "RK4":
         transport_scheme = RK4(domain)
     elif scheme == "Heun":

From 585e66cb75f7a08ab49c9de793b00293789265e2 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Thu, 7 Sep 2023 15:03:27 +0100
Subject: [PATCH 05/13] lint fix

---
 integration-tests/model/test_time_discretisation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py
index a61e43d9e..ebe9993b1 100644
--- a/integration-tests/model/test_time_discretisation.py
+++ b/integration-tests/model/test_time_discretisation.py
@@ -30,9 +30,9 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup):
     elif scheme == "TrapeziumRule":
         transport_scheme = TrapeziumRule(domain)
     elif scheme == "ImplicitMidpoint":
-      transport_scheme = ImplicitMidpoint(domain)
+        transport_scheme = ImplicitMidpoint(domain)
     elif scheme == "QinZhang":
-       transport_scheme = QinZhang(domain)
+        transport_scheme = QinZhang(domain)
     elif scheme == "RK4":
         transport_scheme = RK4(domain)
     elif scheme == "Heun":

From f01d8819c03bee28b7eaeb98676c6d1131e2192a Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Fri, 15 Sep 2023 10:14:57 +0100
Subject: [PATCH 06/13] Adding cached solver list for efficiency

---
 gusto/time_discretisation.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index dce51c963..549dc860e 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -301,7 +301,7 @@ def lhs(self):
     def rhs(self):
         return super().rhs
 
-    def solvers(self, stage):
+    def solver(self, stage):
         residual = self.residual.label_map(
             lambda t: t.has_label(time_derivative),
             map_if_true=drop,
@@ -319,6 +319,13 @@ def solvers(self, stage):
         return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters,
                                           options_prefix=solver_name)
 
+    @cached_property
+    def solvers(self):
+        solvers = []
+        for stage in range(self.nStages):
+            solvers.append(self.solver(stage))
+        return solvers
+
     def solve_stage(self, x0, stage):
         self.x1.assign(x0)
         for i in range(stage):
@@ -331,8 +338,8 @@ def solve_stage(self, x0, stage):
             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
-
-        self.solvers(stage).solve()
+        solver = self.solvers[stage]
+        solver.solve()
 
         self.k[stage].assign(self.x_out)
 

From 7c1bfac26662d30b8232699c19c24e330211389b Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Tue, 19 Sep 2023 11:38:10 +0100
Subject: [PATCH 07/13] Changes in response to review

---
 gusto/time_discretisation.py | 318 ++++++++++++++++++++++++-----------
 1 file changed, 219 insertions(+), 99 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 549dc860e..9906fb781 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -242,44 +242,61 @@ class ImplicitMultistage(TimeDiscretisation):
     A Butcher tableau is formed in the following way for a s-th order
     diagonally implicit scheme:
 
-    c_0 | a_00 a_01  .    a_0s
-    c_1 | a_10 a_11       a_1s
-     .  |   .   .    .    .
-     .  |   .   .    .    .
-    c_s | a_s0 a_s1  .    a_ss
-     -------------------------
-        |  b_1  b_2  ...  b_s
+    c_0 | a_00 a_01  .    a_0s                                                /n
+    c_1 | a_10 a_11       a_1s                                                /n
+     .  |   .   .    .    .                                                   /n
+     .  |   .   .    .    .                                                   /n
+    c_s | a_s0 a_s1  .    a_ss                                                /n
+     -------------------------                                                /n
+        |  b_1  b_2  ...  b_s                                                 /n
 
     The gradient function has no time-dependence, so c elements are not needed.
     A square 'butcher_matrix' is defined in each class that uses
     the ImplicitMultiStage structure,
 
-    [a_00   0   .       0  ]
-    [a_10 a_11  .       0  ]
-    [  .    .   .       .  ]
-    [ b_0  b_1  .       b_s]
+    [a_00   0   .       0  ]                                                  /n
+    [a_10 a_11  .       0  ]                                                  /n
+    [  .    .   .       .  ]                                                  /n
+    [ b_0  b_1  .       b_s]                                                  /n
 
     Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods.
 
     There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
 
-    For an s stage method,
-      At iteration i (from 0 to s-1)
-        An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i)
-        Compute the gradient at the intermediate location, k_i = F(y_i)
+    For each i = 1, s  in an s stage method
+    we have the intermediate solutions:
+    y_i = y^n +  dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i)                    /n
+    We compute the gradient at the intermediate location, k_i = F(y_i)        /n
 
     At the last stage, compute the new solution by:
-    y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i)
+    y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   /n
 
     """
 
-    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None):
+    def __init__(self, domain, butcher_matrix, field_name=None,
+                 solver_parameters=None, limiter=None, options=None,):
+        """
+        Args:
+            domain (:class:`Domain`): the model's domain object, containing the
+                mesh and the compatible function spaces.
+            butcher_matrix (numpy array): A matrix containing the coefficients of
+                a butcher tableau defining a given 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.
+            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
+                recovery method. Defaults to None.
+        """
         super().__init__(domain, field_name=field_name,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        if butcher_matrix is not None:
-            self.butcher_matrix = butcher_matrix
-            self.nStages = int(np.shape(self.butcher_matrix)[1])
+        self.butcher_matrix = butcher_matrix
+        self.nStages = int(np.shape(self.butcher_matrix)[1])
 
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
@@ -464,46 +481,62 @@ class ExplicitMultistage(ExplicitTimeDiscretisation):
 
     A Butcher tableau is formed in the following way for a s-th order explicit scheme:
 
-    c_0 | a_00 a_01  .    a_0s
-    c_1 | a_10 a_11       a_1s
-     .  |   .   .    .    .
-     .  |   .   .    .    .
-    c_s | a_s0 a_s1  .    a_ss
-     -------------------------
-        |  b_1  b_2  ...  b_s
+    c_0 | a_00 a_01  .    a_0s                                                /n
+    c_1 | a_10 a_11       a_1s                                                /n
+     .  |   .   .    .    .                                                   /n
+     .  |   .   .    .    .                                                   /n
+    c_s | a_s0 a_s1  .    a_ss                                                /n
+     -------------------------                                                /n
+        |  b_1  b_2  ...  b_s                                                 /n
 
     The gradient function has no time-dependence, so c elements are not needed.
     A square 'butcher_matrix' is defined in each class that uses
     the ExplicitMultiStage structure,
 
-    [a_10   0   .       0  ]
-    [a_20 a_21  .       0  ]
-    [  .    .   .       .  ]
-    [ b_0  b_1  .       b_s]
+    [a_10   0   .       0  ]                                                  /n
+    [a_20 a_21  .       0  ]                                                  /n
+    [  .    .   .       .  ]                                                  /n
+    [ b_0  b_1  .       b_s]                                                  /n
 
     All upper diagonal a_ij elements are zero for explicit methods. We exclude the first
     row of the butcher tableau from our butcher matrix as the row is always zeros.
 
     There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
 
-    For an s stage method,
-      At iteration i (from 0 to s-1)
-        An intermediate location is computed as y_i = y^n + sum{over j less than i} (dt*a_ij*k_i)
-        Compute the gradient at the intermediate location, k_i = F(y_i)
+    For each i = 1, s  in an s stage method
+    we have the intermediate solutions:
+    y_i = y^n +  dt*(a_i1*k_1 + a_i2*k_2 + ... + a_i{i-1}*k_{i-1})            /n
+    We compute the gradient at the intermediate location, k_i = F(y_i)        /n
 
     At the last stage, compute the new solution by:
-    y^{n+1} = y^n + sum_{j from 0 to s-1} (dt*b_i*k_i)
+    y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   /n
 
     """
 
-    def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
-                 limiter=None, options=None, butcher_matrix=None):
+    def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None,
+                 solver_parameters=None, limiter=None, options=None):
+        """
+        Args:
+            domain (:class:`Domain`): the model's domain object, containing the
+                mesh and the compatible function spaces.
+            butcher_matrix (numpy array): A matrix containing the coefficients of
+                a butcher tableau defining a given 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.
+            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
+                recovery method. Defaults to None.
+        """
         super().__init__(domain, field_name=field_name, subcycles=subcycles,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        if butcher_matrix is not None:
-            self.butcher_matrix = butcher_matrix
-            self.nbutcher = int(np.shape(self.butcher_matrix)[0])
+        self.butcher_matrix = butcher_matrix
+        self.nbutcher = int(np.shape(self.butcher_matrix)[0])
 
     @property
     def nStages(self):
@@ -584,16 +617,31 @@ class ForwardEuler(ExplicitMultistage):
     Implements the forward Euler timestepping scheme.
 
     The forward Euler method for operator F is the most simple explicit scheme:
-    k0 = F[y^n]
-    y^(n+1) = y^n + dt*k0
+    k0 = F[y^n]                                                               /n
+    y^(n+1) = y^n + dt*k0                                                     /n
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
-                 limiter=None, options=None, butcher_matrix=None):
-        super().__init__(domain, field_name=field_name, subcycles=subcycles,
+                 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.
+            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
+                recovery method. Defaults to None.
+        """
+        butcher_matrix = np.array([1.]).reshape(1, 1)
+        super().__init__(domain, butcher_matrix=butcher_matrix,
+                         field_name=field_name, subcycles=subcycles,
                          solver_parameters=solver_parameters,
-                         limiter=limiter, options=options, butcher_matrix=butcher_matrix)
-        self.butcher_matrix = np.array([1.]).reshape(1, 1)
-        self.nbutcher = int(np.shape(self.butcher_matrix)[0])
+                         limiter=limiter, options=options)
 
 
 class SSPRK3(ExplicitMultistage):
@@ -601,18 +649,32 @@ class SSPRK3(ExplicitMultistage):
     Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method
     for solving ∂y/∂t = F(y). It can be written as:
 
-    k0 = F[y^n]
-    k1 = F[y^n + dt*k1]
-    k2 = F[y^n + (1/4)*dt*(k0+k1)]
-    y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2)
+    k0 = F[y^n]                                                               /n
+    k1 = F[y^n + dt*k1]                                                       /n
+    k2 = F[y^n + (1/4)*dt*(k0+k1)]                                            /n
+    y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2)                                 /n
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
-                 limiter=None, options=None, butcher_matrix=None):
-        super().__init__(domain, field_name=field_name, subcycles=subcycles,
-                         solver_parameters=solver_parameters,
-                         limiter=limiter, options=options, butcher_matrix=butcher_matrix)
-        self.butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]])
-        self.nbutcher = int(np.shape(self.butcher_matrix)[0])
+                 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.
+            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
+                recovery method. Defaults to None.
+        """
+        butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]])
+        super().__init__(domain, butcher_matrix=butcher_matrix, field_name=field_name,
+                         subcycles=subcycles, solver_parameters=solver_parameters,
+                         limiter=limiter, options=options)
 
 
 class RK4(ExplicitMultistage):
@@ -622,21 +684,35 @@ class RK4(ExplicitMultistage):
     The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be
     written as:
 
-    k0 = F[y^n]
-    k1 = F[y^n + 1/2*dt*k1]
-    k2 = F[y^n + 1/2*dt*k2]
-    k3 = F[y^n + dt*k3]
-    y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3)
+    k0 = F[y^n]                                                               /n
+    k1 = F[y^n + 1/2*dt*k1]                                                   /n
+    k2 = F[y^n + 1/2*dt*k2]                                                   /n
+    k3 = F[y^n + dt*k3]                                                       /n
+    y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3)                      /n
 
-    where superscripts indicate the time-level.
+    where superscripts indicate the time-level.                               /n
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
-                 limiter=None, options=None, butcher_matrix=None):
-        super().__init__(domain, field_name=field_name, subcycles=subcycles,
-                         solver_parameters=solver_parameters,
-                         limiter=limiter, options=options, butcher_matrix=butcher_matrix)
-        self.butcher_matrix = np.array([[0.5, 0., 0., 0.], [0., 0.5, 0., 0.], [0., 0., 1., 0.], [1./6., 1./3., 1./3., 1./6.]])
-        self.nbutcher = int(np.shape(self.butcher_matrix)[0])
+                 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.
+            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
+                recovery method. Defaults to None.
+        """
+        butcher_matrix = np.array([[0.5, 0., 0., 0.], [0., 0.5, 0., 0.], [0., 0., 1., 0.], [1./6., 1./3., 1./3., 1./6.]])
+        super().__init__(domain, butcher_matrix=butcher_matrix, field_name=field_name,
+                         subcycles=subcycles, solver_parameters=solver_parameters,
+                         limiter=limiter, options=options)
 
 
 class Heun(ExplicitMultistage):
@@ -646,19 +722,33 @@ class Heun(ExplicitMultistage):
     The 2-stage Runge-Kutta scheme known as Heun's method,for solving
     ∂y/∂t = F(y). It can be written as:
 
-    y_1 = F[y^n]
-    y^(n+1) = (1/2)y^n + (1/2)F[y_1]
+    y_1 = F[y^n]                                                              /n
+    y^(n+1) = (1/2)y^n + (1/2)F[y_1]                                          /n
 
     where superscripts indicate the time-level and subscripts indicate the stage
     number.
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
-                 limiter=None, options=None, butcher_matrix=None):
-        super().__init__(domain, field_name,
-                         solver_parameters=solver_parameters,
+                 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.
+            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
+                recovery method. Defaults to None.
+        """
+        butcher_matrix = np.array([[1., 0.], [0.5, 0.5]])
+        super().__init__(domain, butcher_matrix=butcher_matrix,
+                         field_name=field_name, solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        self.butcher_matrix = np.array([[1., 0.], [0.5, 0.5]])
-        self.nbutcher = np.shape(self.butcher_matrix)[0]
 
 
 class BackwardEuler(TimeDiscretisation):
@@ -741,7 +831,7 @@ class ThetaMethod(TimeDiscretisation):
     be thought as a generalised trapezium rule.
 
     The theta implicit-explicit timestepping method for operator F is written as
-    y^(n+1) = y^n + dt*(1-theta)*F[y^n] + dt*theta*F[y^(n+1)]
+    y^(n+1) = y^n + dt*(1-theta)*F[y^n] + dt*theta*F[y^(n+1)]                 /n
     for off-centring parameter theta.
     """
 
@@ -889,7 +979,7 @@ class BDF2(MultilevelTimeDiscretisation):
     Implements the implicit multistep BDF2 timestepping method
 
     The BDF2 timestepping method for operator F is written as
-    y^(n+1) = (4/3)*y^n - (1/3)*y^(n-1) + (2/3)*dt*F[y^(n+1)]
+    y^(n+1) = (4/3)*y^n - (1/3)*y^(n-1) + (2/3)*dt*F[y^(n+1)]                 /n
     """
 
     @property
@@ -988,9 +1078,9 @@ class TR_BDF2(TimeDiscretisation):
     trapezoidal stage (TR) followed by a second order backwards difference stage (BDF2).
 
     The TR-BDF2 time stepping method for operator F is written as
-    y^(n+g) = y^n + dt*g/2*F[y^n] + dt*g/2*F[y^(n+g)] (TR stage)
-    y^(n+1) = 1/(g(2-g))*y^(n+g) - (1-g)**2/(g(2-g))*y^(n) + (1-g)/(2-g)*dt*F[y^(n+1)] (BDF2 stage)
-    for an off-centring parameter g (gamma).
+    y^(n+g) = y^n + dt*g/2*F[y^n] + dt*g/2*F[y^(n+g)] (TR stage)                                    /n
+    y^(n+1) = 1/(g(2-g))*y^(n+g) - (1-g)**2/(g(2-g))*y^(n) + (1-g)/(2-g)*dt*F[y^(n+1)] (BDF2 stage) /n
+    for an off-centring parameter g (gamma).                                                        /n
     """
     def __init__(self, domain, gamma, field_name=None,
                  solver_parameters=None, options=None):
@@ -1193,7 +1283,7 @@ class AdamsBashforth(MultilevelTimeDiscretisation):
     Implements the explicit multistep Adams-Bashforth timestepping method of general order up to 5
 
     The general AB timestepping method for operator F is written as
-    y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)])
+    y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) /n
     """
     def __init__(self, domain, order, field_name=None,
                  solver_parameters=None, options=None):
@@ -1323,7 +1413,7 @@ class AdamsMoulton(MultilevelTimeDiscretisation):
     Implements the implicit multistep Adams-Moulton timestepping method of general order up to 5
 
     The general AM timestepping method for operator F is written as
-    y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)])
+    y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) /n
     """
     def __init__(self, domain, order, field_name=None,
                  solver_parameters=None, options=None):
@@ -1472,17 +1562,32 @@ class ImplicitMidpoint(ImplicitMultistage):
     Implements the Implicit Midpoint method as a 1-stage Runge Kutta method.
 
     The method, for solving
-    ∂y/∂t = F(y), can be written as:
+    ∂y/∂t = F(y), can be written as:                                          /n
 
-    k0 = F[y^n + 0.5*dt*k0]
-    y^(n+1) = y^n + dt*k0
+    k0 = F[y^n + 0.5*dt*k0]                                                   /n
+    y^(n+1) = y^n + dt*k0                                                     /n
     """
-    def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None):
-        super().__init__(domain, field_name,
+    def __init__(self, domain, field_name=None, 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.
+            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
+                recovery method. Defaults to None.
+        """
+        butcher_matrix = np.array([[0.5], [1.]])
+        super().__init__(domain, field_name, butcher_matrix=butcher_matrix,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        self.butcher_matrix = np.array([[0.5], [1.]])
-        self.nStages = int(np.shape(self.butcher_matrix)[1])
 
 
 class QinZhang(ImplicitMultistage):
@@ -1490,15 +1595,30 @@ class QinZhang(ImplicitMultistage):
     Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method.
 
     The method, for solving
-    ∂y/∂t = F(y), can be written as:
+    ∂y/∂t = F(y), can be written as:                                          /n
 
-    k0 = F[y^n + 0.25*dt*k0]
-    k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1]
-    y^(n+1) = y^n + 0.5*dt*(k0 + k1)
+    k0 = F[y^n + 0.25*dt*k0]                                                  /n
+    k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1]                                      /n
+    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, butcher_matrix=None):
-        super().__init__(domain, field_name,
+    def __init__(self, domain, field_name=None, 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.
+            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
+                recovery method. Defaults to None.
+        """
+        butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]])
+        super().__init__(domain, field_name, butcher_matrix=butcher_matrix,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
-        self.butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]])
-        self.nStages = int(np.shape(self.butcher_matrix)[1])

From c62fe1f62fee4bdd9a333f73ec834a52215733ff Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Tue, 19 Sep 2023 11:55:42 +0100
Subject: [PATCH 08/13] correct minor error

---
 gusto/time_discretisation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 9906fb781..71f3b422f 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -1585,7 +1585,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 recovery method. Defaults to None.
         """
         butcher_matrix = np.array([[0.5], [1.]])
-        super().__init__(domain, field_name, butcher_matrix=butcher_matrix,
+        super().__init__(domain, butcher_matrix, field_name,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
 
@@ -1619,6 +1619,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 recovery method. Defaults to None.
         """
         butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]])
-        super().__init__(domain, field_name, butcher_matrix=butcher_matrix,
+        super().__init__(domain, butcher_matrix, field_name,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)

From b005650fb376503dde26bd8b2b54065a88eef651 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Tue, 19 Sep 2023 13:52:34 +0100
Subject: [PATCH 09/13] subcycles added to Heuns, and comments added

---
 gusto/time_discretisation.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 71f3b422f..53c533479 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -523,6 +523,8 @@ def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None,
                 a butcher tableau defining a given Runge Kutta time discretisation.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
+            subcycles (int, optional): the number of sub-steps to perform.
+                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
@@ -628,6 +630,8 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
+            subcycles (int, optional): the number of sub-steps to perform.
+                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
@@ -662,6 +666,8 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
+            subcycles (int, optional): the number of sub-steps to perform.
+                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
@@ -700,6 +706,8 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
+            subcycles (int, optional): the number of sub-steps to perform.
+                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
@@ -736,6 +744,8 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
+            subcycles (int, optional): the number of sub-steps to perform.
+                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
@@ -747,7 +757,8 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No
         """
         butcher_matrix = np.array([[1., 0.], [0.5, 0.5]])
         super().__init__(domain, butcher_matrix=butcher_matrix,
-                         field_name=field_name, solver_parameters=solver_parameters,
+                         field_name=field_name, subcycles=subcycles,
+                         solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
 
 

From 5bb0b34e63742e71311c94377397e259cf736237 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Thu, 21 Sep 2023 11:05:09 +0100
Subject: [PATCH 10/13] Tidy

---
 gusto/time_discretisation.py | 172 +++++++++++++++++------------------
 1 file changed, 86 insertions(+), 86 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 53c533479..89ae94113 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -240,36 +240,34 @@ class ImplicitMultistage(TimeDiscretisation):
     methods based on its Butcher tableau.
 
     A Butcher tableau is formed in the following way for a s-th order
-    diagonally implicit scheme:
-
-    c_0 | a_00 a_01  .    a_0s                                                /n
-    c_1 | a_10 a_11       a_1s                                                /n
-     .  |   .   .    .    .                                                   /n
-     .  |   .   .    .    .                                                   /n
-    c_s | a_s0 a_s1  .    a_ss                                                /n
-     -------------------------                                                /n
-        |  b_1  b_2  ...  b_s                                                 /n
-
+    diagonally implicit scheme:                                               \n
+    c_0 | a_00 a_01  .    a_0s                                                \n
+    c_1 | a_10 a_11       a_1s                                                \n
+     .  |   .   .    .    .                                                   \n
+     .  |   .   .    .    .                                                   \n
+    c_s | a_s0 a_s1  .    a_ss                                                \n
+     -------------------------                                                \n
+        |  b_1  b_2  ...  b_s                                                 \n
+    
     The gradient function has no time-dependence, so c elements are not needed.
     A square 'butcher_matrix' is defined in each class that uses
-    the ImplicitMultiStage structure,
-
-    [a_00   0   .       0  ]                                                  /n
-    [a_10 a_11  .       0  ]                                                  /n
-    [  .    .   .       .  ]                                                  /n
-    [ b_0  b_1  .       b_s]                                                  /n
+    the ImplicitMultiStage structure,                                         \n
+    [a_00   0   .       0  ]                                                  \n
+    [a_10 a_11  .       0  ]                                                  \n
+    [  .    .   .       .  ]                                                  \n
+    [ b_0  b_1  .       b_s]                                                  \n
 
     Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods.
 
     There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
 
     For each i = 1, s  in an s stage method
-    we have the intermediate solutions:
-    y_i = y^n +  dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i)                    /n
-    We compute the gradient at the intermediate location, k_i = F(y_i)        /n
+    we have the intermediate solutions:                                       \n
+    y_i = y^n +  dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i)                    \n
+    We compute the gradient at the intermediate location, k_i = F(y_i)        \n
 
-    At the last stage, compute the new solution by:
-    y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   /n
+    At the last stage, compute the new solution by:                           \n
+    y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   \n
 
     """
 
@@ -479,24 +477,24 @@ class ExplicitMultistage(ExplicitTimeDiscretisation):
     A class for implementing general explicit multistage (Runge-Kutta)
     methods based on its Butcher tableau.
 
-    A Butcher tableau is formed in the following way for a s-th order explicit scheme:
+    A Butcher tableau is formed in the following way for a s-th order explicit scheme: \n
 
-    c_0 | a_00 a_01  .    a_0s                                                /n
-    c_1 | a_10 a_11       a_1s                                                /n
-     .  |   .   .    .    .                                                   /n
-     .  |   .   .    .    .                                                   /n
-    c_s | a_s0 a_s1  .    a_ss                                                /n
-     -------------------------                                                /n
-        |  b_1  b_2  ...  b_s                                                 /n
+    c_0 | a_00 a_01  .    a_0s                                                \n
+    c_1 | a_10 a_11       a_1s                                                \n
+     .  |   .   .    .    .                                                   \n
+     .  |   .   .    .    .                                                   \n
+    c_s | a_s0 a_s1  .    a_ss                                                \n
+     -------------------------                                                \n
+        |  b_1  b_2  ...  b_s                                                 \n
 
     The gradient function has no time-dependence, so c elements are not needed.
     A square 'butcher_matrix' is defined in each class that uses
-    the ExplicitMultiStage structure,
+    the ExplicitMultiStage structure, \n
 
-    [a_10   0   .       0  ]                                                  /n
-    [a_20 a_21  .       0  ]                                                  /n
-    [  .    .   .       .  ]                                                  /n
-    [ b_0  b_1  .       b_s]                                                  /n
+    [a_10   0   .       0  ]                                                  \n
+    [a_20 a_21  .       0  ]                                                  \n
+    [  .    .   .       .  ]                                                  \n
+    [ b_0  b_1  .       b_s]                                                  \n
 
     All upper diagonal a_ij elements are zero for explicit methods. We exclude the first
     row of the butcher tableau from our butcher matrix as the row is always zeros.
@@ -504,12 +502,12 @@ class ExplicitMultistage(ExplicitTimeDiscretisation):
     There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
 
     For each i = 1, s  in an s stage method
-    we have the intermediate solutions:
-    y_i = y^n +  dt*(a_i1*k_1 + a_i2*k_2 + ... + a_i{i-1}*k_{i-1})            /n
-    We compute the gradient at the intermediate location, k_i = F(y_i)        /n
+    we have the intermediate solutions:                                       \n
+    y_i = y^n +  dt*(a_i1*k_1 + a_i2*k_2 + ... + a_i{i-1}*k_{i-1})            \n
+    We compute the gradient at the intermediate location, k_i = F(y_i)        \n
 
     At the last stage, compute the new solution by:
-    y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   /n
+    y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   \n
 
     """
 
@@ -618,9 +616,9 @@ class ForwardEuler(ExplicitMultistage):
     """
     Implements the forward Euler timestepping scheme.
 
-    The forward Euler method for operator F is the most simple explicit scheme:
-    k0 = F[y^n]                                                               /n
-    y^(n+1) = y^n + dt*k0                                                     /n
+    The forward Euler method for operator F is the most simple explicit scheme: \n
+    k0 = F[y^n]                                                                 \n
+    y^(n+1) = y^n + dt*k0                                                       \n
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
                  limiter=None, options=None):
@@ -651,12 +649,12 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No
 class SSPRK3(ExplicitMultistage):
     u"""
     Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method
-    for solving ∂y/∂t = F(y). It can be written as:
+    for solving ∂y/∂t = F(y). It can be written as:                           \n
 
-    k0 = F[y^n]                                                               /n
-    k1 = F[y^n + dt*k1]                                                       /n
-    k2 = F[y^n + (1/4)*dt*(k0+k1)]                                            /n
-    y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2)                                 /n
+    k0 = F[y^n]                                                               \n
+    k1 = F[y^n + dt*k1]                                                       \n
+    k2 = F[y^n + (1/4)*dt*(k0+k1)]                                            \n
+    y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2)                                 \n
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
                  limiter=None, options=None):
@@ -688,15 +686,15 @@ class RK4(ExplicitMultistage):
     Implements the classic 4-stage Runge-Kutta method.
 
     The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be
-    written as:
+    written as:                                                               \n
 
-    k0 = F[y^n]                                                               /n
-    k1 = F[y^n + 1/2*dt*k1]                                                   /n
-    k2 = F[y^n + 1/2*dt*k2]                                                   /n
-    k3 = F[y^n + dt*k3]                                                       /n
-    y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3)                      /n
+    k0 = F[y^n]                                                               \n
+    k1 = F[y^n + 1/2*dt*k1]                                                   \n
+    k2 = F[y^n + 1/2*dt*k2]                                                   \n
+    k3 = F[y^n + dt*k3]                                                       \n
+    y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3)                      \n
 
-    where superscripts indicate the time-level.                               /n
+    where superscripts indicate the time-level.                               \n
     """
     def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None,
                  limiter=None, options=None):
@@ -728,10 +726,10 @@ class Heun(ExplicitMultistage):
     Implements Heun's method.
 
     The 2-stage Runge-Kutta scheme known as Heun's method,for solving
-    ∂y/∂t = F(y). It can be written as:
+    ∂y/∂t = F(y). It can be written as:                                       \n
 
-    y_1 = F[y^n]                                                              /n
-    y^(n+1) = (1/2)y^n + (1/2)F[y_1]                                          /n
+    y_1 = F[y^n]                                                              \n
+    y^(n+1) = (1/2)y^n + (1/2)F[y_1]                                          \n
 
     where superscripts indicate the time-level and subscripts indicate the stage
     number.
@@ -766,8 +764,8 @@ class BackwardEuler(TimeDiscretisation):
     """
     Implements the backward Euler timestepping scheme.
 
-    The backward Euler method for operator F is the most simple implicit scheme:
-    y^(n+1) = y^n + dt*F[y^(n+1)].
+    The backward Euler method for operator F is the most simple implicit scheme: \n
+    y^(n+1) = y^n + dt*F[y^(n+1)].                                               \n
     """
     def __init__(self, domain, field_name=None, solver_parameters=None,
                  limiter=None, options=None):
@@ -841,9 +839,9 @@ class ThetaMethod(TimeDiscretisation):
     Implements the theta implicit-explicit timestepping method, which can
     be thought as a generalised trapezium rule.
 
-    The theta implicit-explicit timestepping method for operator F is written as
-    y^(n+1) = y^n + dt*(1-theta)*F[y^n] + dt*theta*F[y^(n+1)]                 /n
-    for off-centring parameter theta.
+    The theta implicit-explicit timestepping method for operator F is written as: \n
+    y^(n+1) = y^n + dt*(1-theta)*F[y^n] + dt*theta*F[y^(n+1)]                     \n
+    for off-centring parameter theta.                                             \n
     """
 
     def __init__(self, domain, theta, field_name=None,
@@ -923,9 +921,9 @@ class TrapeziumRule(ThetaMethod):
     Implements the trapezium rule timestepping method, also commonly known as
     Crank Nicholson.
 
-    The trapezium rule timestepping method for operator F is written as
-    y^(n+1) = y^n + dt/2*F[y^n] + dt/2*F[y^(n+1)].
-    It is equivalent to the "theta" method with theta = 1/2.
+    The trapezium rule timestepping method for operator F is written as:      \n
+    y^(n+1) = y^n + dt/2*F[y^n] + dt/2*F[y^(n+1)].                            \n
+    It is equivalent to the "theta" method with theta = 1/2.                  \n
     """
 
     def __init__(self, domain, field_name=None, solver_parameters=None,
@@ -987,10 +985,10 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
 class BDF2(MultilevelTimeDiscretisation):
     """
-    Implements the implicit multistep BDF2 timestepping method
+    Implements the implicit multistep BDF2 timestepping method.
 
-    The BDF2 timestepping method for operator F is written as
-    y^(n+1) = (4/3)*y^n - (1/3)*y^(n-1) + (2/3)*dt*F[y^(n+1)]                 /n
+    The BDF2 timestepping method for operator F is written as:                \n
+    y^(n+1) = (4/3)*y^n - (1/3)*y^(n-1) + (2/3)*dt*F[y^(n+1)]                 \n
     """
 
     @property
@@ -1088,10 +1086,10 @@ class TR_BDF2(TimeDiscretisation):
     Implements the two stage implicit TR-BDF2 time stepping method, with a
     trapezoidal stage (TR) followed by a second order backwards difference stage (BDF2).
 
-    The TR-BDF2 time stepping method for operator F is written as
-    y^(n+g) = y^n + dt*g/2*F[y^n] + dt*g/2*F[y^(n+g)] (TR stage)                                    /n
-    y^(n+1) = 1/(g(2-g))*y^(n+g) - (1-g)**2/(g(2-g))*y^(n) + (1-g)/(2-g)*dt*F[y^(n+1)] (BDF2 stage) /n
-    for an off-centring parameter g (gamma).                                                        /n
+    The TR-BDF2 time stepping method for operator F is written as:                                  \n
+    y^(n+g) = y^n + dt*g/2*F[y^n] + dt*g/2*F[y^(n+g)] (TR stage)                                    \n
+    y^(n+1) = 1/(g(2-g))*y^(n+g) - (1-g)**2/(g(2-g))*y^(n) + (1-g)/(2-g)*dt*F[y^(n+1)] (BDF2 stage) \n
+    for an off-centring parameter g (gamma).                                                        \n
     """
     def __init__(self, domain, gamma, field_name=None,
                  solver_parameters=None, options=None):
@@ -1216,8 +1214,8 @@ class Leapfrog(MultilevelTimeDiscretisation):
     """
     Implements the multistep Leapfrog timestepping method.
 
-    The Leapfrog timestepping method for operator F is written as
-    y^(n+1) = y^(n-1)  + 2*dt*F[y^n]
+    The Leapfrog timestepping method for operator F is written as:            \n
+    y^(n+1) = y^(n-1)  + 2*dt*F[y^n]                                          \n
     """
     @property
     def nlevels(self):
@@ -1291,10 +1289,11 @@ def apply(self, x_out, *x_in):
 
 class AdamsBashforth(MultilevelTimeDiscretisation):
     """
-    Implements the explicit multistep Adams-Bashforth timestepping method of general order up to 5
+    Implements the explicit multistep Adams-Bashforth timestepping 
+    method of general order up to 5.
 
-    The general AB timestepping method for operator F is written as
-    y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) /n
+    The general AB timestepping method for operator F is written as:                                      \n
+    y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) \n
     """
     def __init__(self, domain, order, field_name=None,
                  solver_parameters=None, options=None):
@@ -1421,10 +1420,11 @@ def apply(self, x_out, *x_in):
 
 class AdamsMoulton(MultilevelTimeDiscretisation):
     """
-    Implements the implicit multistep Adams-Moulton timestepping method of general order up to 5
+    Implements the implicit multistep Adams-Moulton 
+    timestepping method of general order up to 5
 
-    The general AM timestepping method for operator F is written as
-    y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) /n
+    The general AM timestepping method for operator F is written as                      \n
+    y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) \n
     """
     def __init__(self, domain, order, field_name=None,
                  solver_parameters=None, options=None):
@@ -1573,10 +1573,10 @@ class ImplicitMidpoint(ImplicitMultistage):
     Implements the Implicit Midpoint method as a 1-stage Runge Kutta method.
 
     The method, for solving
-    ∂y/∂t = F(y), can be written as:                                          /n
+    ∂y/∂t = F(y), can be written as:                                          \n
 
-    k0 = F[y^n + 0.5*dt*k0]                                                   /n
-    y^(n+1) = y^n + dt*k0                                                     /n
+    k0 = F[y^n + 0.5*dt*k0]                                                   \n
+    y^(n+1) = y^n + dt*k0                                                     \n
     """
     def __init__(self, domain, field_name=None, solver_parameters=None,
                  limiter=None, options=None):
@@ -1606,11 +1606,11 @@ class QinZhang(ImplicitMultistage):
     Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method.
 
     The method, for solving
-    ∂y/∂t = F(y), can be written as:                                          /n
+    ∂y/∂t = F(y), can be written as:                                          \n
 
-    k0 = F[y^n + 0.25*dt*k0]                                                  /n
-    k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1]                                      /n
-    y^(n+1) = y^n + 0.5*dt*(k0 + k1)                                          /n
+    k0 = F[y^n + 0.25*dt*k0]                                                  \n
+    k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1]                                      \n
+    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):

From ac4c08b9afe6bc2ea3017e9a199e8896bb7180d8 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Tue, 17 Oct 2023 16:13:38 +0100
Subject: [PATCH 11/13] moving butcher matrix comment

---
 gusto/time_discretisation.py | 28 ++++++++++------------------
 1 file changed, 10 insertions(+), 18 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 89ae94113..8899b3f67 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -239,24 +239,6 @@ class ImplicitMultistage(TimeDiscretisation):
     A class for implementing general diagonally implicit multistage (Runge-Kutta)
     methods based on its Butcher tableau.
 
-    A Butcher tableau is formed in the following way for a s-th order
-    diagonally implicit scheme:                                               \n
-    c_0 | a_00 a_01  .    a_0s                                                \n
-    c_1 | a_10 a_11       a_1s                                                \n
-     .  |   .   .    .    .                                                   \n
-     .  |   .   .    .    .                                                   \n
-    c_s | a_s0 a_s1  .    a_ss                                                \n
-     -------------------------                                                \n
-        |  b_1  b_2  ...  b_s                                                 \n
-    
-    The gradient function has no time-dependence, so c elements are not needed.
-    A square 'butcher_matrix' is defined in each class that uses
-    the ImplicitMultiStage structure,                                         \n
-    [a_00   0   .       0  ]                                                  \n
-    [a_10 a_11  .       0  ]                                                  \n
-    [  .    .   .       .  ]                                                  \n
-    [ b_0  b_1  .       b_s]                                                  \n
-
     Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods.
 
     There are three steps to move from the current solution, y^n, to the new one, y^{n+1}
@@ -279,6 +261,11 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                 mesh and the compatible function spaces.
             butcher_matrix (numpy array): A matrix containing the coefficients of
                 a butcher tableau defining a given Runge Kutta time discretisation.
+                The matrix is of the form:
+                [a_00   0   .       0  ]
+                [a_10 a_11  .       0  ]
+                [  .    .   .       .  ]
+                [ b_0  b_1  .       b_s]
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
             solver_parameters (dict, optional): dictionary of parameters to
@@ -519,6 +506,11 @@ def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None,
                 mesh and the compatible function spaces.
             butcher_matrix (numpy array): A matrix containing the coefficients of
                 a butcher tableau defining a given Runge Kutta time discretisation.
+                The matrix is of the form:
+                [a_10   0   .       0  ]
+                [a_20 a_21  .       0  ]
+                [  .    .   .       .  ]
+                [ b_0  b_1  .       b_s]
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
             subcycles (int, optional): the number of sub-steps to perform.

From 7fed4c747baac873d954977be94f39ca04f8a9a1 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Tue, 17 Oct 2023 16:15:06 +0100
Subject: [PATCH 12/13] lint fix

---
 gusto/time_discretisation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 8899b3f67..c5b14aed5 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -1281,7 +1281,7 @@ def apply(self, x_out, *x_in):
 
 class AdamsBashforth(MultilevelTimeDiscretisation):
     """
-    Implements the explicit multistep Adams-Bashforth timestepping 
+    Implements the explicit multistep Adams-Bashforth timestepping
     method of general order up to 5.
 
     The general AB timestepping method for operator F is written as:                                      \n
@@ -1412,7 +1412,7 @@ def apply(self, x_out, *x_in):
 
 class AdamsMoulton(MultilevelTimeDiscretisation):
     """
-    Implements the implicit multistep Adams-Moulton 
+    Implements the implicit multistep Adams-Moulton
     timestepping method of general order up to 5
 
     The general AM timestepping method for operator F is written as                      \n

From 8ceeb01e2e1224f136dae78a3ee75206f4443932 Mon Sep 17 00:00:00 2001
From: atb1995 <alex.brown1995@hotmail.co.uk>
Date: Tue, 17 Oct 2023 16:37:56 +0100
Subject: [PATCH 13/13] tidy of comments

---
 gusto/time_discretisation.py | 65 +++++++++++++++++++++---------------
 1 file changed, 38 insertions(+), 27 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index c5b14aed5..1e1ff41b1 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -252,6 +252,25 @@ class ImplicitMultistage(TimeDiscretisation):
     y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   \n
 
     """
+    # ---------------------------------------------------------------------------
+    # Butcher tableau for a s-th order
+    # diagonally implicit scheme:
+    #  c_0 | a_00  0    .     0
+    #  c_1 | a_10 a_11  .     0
+    #   .  |   .   .    .     .
+    #   .  |   .   .    .     .
+    #  c_s | a_s0 a_s1  .    a_ss
+    #   -------------------------
+    #      |  b_1  b_2  ...  b_s
+    #
+    #
+    # The corresponding square 'butcher_matrix' is:
+    #
+    #  [a_00   0   .       0  ]
+    #  [a_10 a_11  .       0  ]
+    #  [  .    .   .       .  ]
+    #  [ b_0  b_1  .       b_s]
+    # ---------------------------------------------------------------------------
 
     def __init__(self, domain, butcher_matrix, field_name=None,
                  solver_parameters=None, limiter=None, options=None,):
@@ -261,11 +280,6 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                 mesh and the compatible function spaces.
             butcher_matrix (numpy array): A matrix containing the coefficients of
                 a butcher tableau defining a given Runge Kutta time discretisation.
-                The matrix is of the form:
-                [a_00   0   .       0  ]
-                [a_10 a_11  .       0  ]
-                [  .    .   .       .  ]
-                [ b_0  b_1  .       b_s]
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
             solver_parameters (dict, optional): dictionary of parameters to
@@ -466,23 +480,6 @@ class ExplicitMultistage(ExplicitTimeDiscretisation):
 
     A Butcher tableau is formed in the following way for a s-th order explicit scheme: \n
 
-    c_0 | a_00 a_01  .    a_0s                                                \n
-    c_1 | a_10 a_11       a_1s                                                \n
-     .  |   .   .    .    .                                                   \n
-     .  |   .   .    .    .                                                   \n
-    c_s | a_s0 a_s1  .    a_ss                                                \n
-     -------------------------                                                \n
-        |  b_1  b_2  ...  b_s                                                 \n
-
-    The gradient function has no time-dependence, so c elements are not needed.
-    A square 'butcher_matrix' is defined in each class that uses
-    the ExplicitMultiStage structure, \n
-
-    [a_10   0   .       0  ]                                                  \n
-    [a_20 a_21  .       0  ]                                                  \n
-    [  .    .   .       .  ]                                                  \n
-    [ b_0  b_1  .       b_s]                                                  \n
-
     All upper diagonal a_ij elements are zero for explicit methods. We exclude the first
     row of the butcher tableau from our butcher matrix as the row is always zeros.
 
@@ -497,6 +494,25 @@ class ExplicitMultistage(ExplicitTimeDiscretisation):
     y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s)                   \n
 
     """
+    # ---------------------------------------------------------------------------
+    # Butcher tableau for a s-th order
+    # explicit scheme:
+    #  c_0 |   0   0    .    0
+    #  c_1 | a_10  0    .    0
+    #   .  |   .   .    .    .
+    #   .  |   .   .    .    .
+    #  c_s | a_s0 a_s1  .    0
+    #   -------------------------
+    #      |  b_1  b_2  ...  b_s
+    #
+    #
+    # The corresponding square 'butcher_matrix' is:
+    #
+    #  [a_10   0   .       0  ]
+    #  [a_20  a_21 .       0  ]
+    #  [  .    .   .       .  ]
+    #  [ b_0  b_1  .       b_s]
+    # ---------------------------------------------------------------------------
 
     def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None,
                  solver_parameters=None, limiter=None, options=None):
@@ -506,11 +522,6 @@ def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None,
                 mesh and the compatible function spaces.
             butcher_matrix (numpy array): A matrix containing the coefficients of
                 a butcher tableau defining a given Runge Kutta time discretisation.
-                The matrix is of the form:
-                [a_10   0   .       0  ]
-                [a_20 a_21  .       0  ]
-                [  .    .   .       .  ]
-                [ b_0  b_1  .       b_s]
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
             subcycles (int, optional): the number of sub-steps to perform.