From 221449551d97595a7c0e3befa364497d49f597e2 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Wed, 15 Nov 2023 15:47:15 +0000
Subject: [PATCH 01/38] add gravity label

---
 gusto/equations.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/gusto/equations.py b/gusto/equations.py
index bd61e795b..dfd174aad 100644
--- a/gusto/equations.py
+++ b/gusto/equations.py
@@ -985,7 +985,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None,
         # -------------------------------------------------------------------- #
         # Gravitational Term
         # -------------------------------------------------------------------- #
-        gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), 'u'), self.X)
+        gravity_form = name_label(subject(prognostic(Term(g*inner(domain.k, w)*dx),
+                                                     'u'), self.X), "gravity")
 
         residual = (mass_form + adv_form + pressure_gradient_form + gravity_form)
 

From 573ad527e91b2f67e470f11451748c76d0e74183 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 20 Dec 2023 02:59:59 +0000
Subject: [PATCH 02/38] started MixedOptions and an integration test for this

---
 gusto/time_discretisation.py                  |   3 +
 gusto/wrappers.py                             |  44 ++-
 .../transport/test_mixed_fs_options.py        | 313 ++++++++++++++++++
 3 files changed, 359 insertions(+), 1 deletion(-)
 create mode 100644 integration-tests/transport/test_mixed_fs_options.py

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 79c8af94e..869b61931 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -87,6 +87,9 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         self.courant_max = None
 
         if options is not None:
+            if options.wrapper_type = Mixed:
+                pass
+        else:
             self.wrapper_name = options.name
             if self.wrapper_name == "embedded_dg":
                 self.wrapper = EmbeddedDGWrapper(self, options)
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 77d46979a..1fd695685 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -15,7 +15,7 @@
 from gusto.labels import transporting_velocity
 import ufl
 
-__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper"]
+__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedOptions"]
 
 
 class Wrapper(object, metaclass=ABCMeta):
@@ -33,6 +33,7 @@ def __init__(self, time_discretisation, wrapper_options):
         self.time_discretisation = time_discretisation
         self.options = wrapper_options
         self.solver_parameters = None
+        self.field_name = None
 
     @abstractmethod
     def setup(self):
@@ -361,3 +362,44 @@ def label_terms(self, residual):
         new_residual = transporting_velocity.update_value(new_residual, self.transporting_velocity)
 
         return new_residual
+        
+
+class MixedOptions(object):
+    """
+    An object to hold a dictionary with different options for different
+    tracers. This means that different tracers can be solved simultaneously
+    using a CoupledTransportEquation, whilst being in different spaces
+    and needing different implementation options. 
+    """
+    
+    def __init__(self, equation, suboptions):
+        """
+        Args:
+            equation (:class: `PrognosticEquationSet`): the prognostic equation(s)
+            sublimiters (dict): A dictionary holding limiters defined for individual prognostic variables
+        Raises:
+            ValueError: If an option is defined for a field that is not in the prognostic variable set
+        """
+    
+    self.suboptions = suboptions
+    self.wrapper_type = Mixed
+    
+    for field, suboption in suboptions:
+        # Check that the field is in the prognostic variable set:
+            if field not in equation.field_names:
+                raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
+            else:
+                # check that a valid wrapper has been given
+                
+                
+                self.suboptions[field].idx = equation.field_names.index(field)
+    
+    
+    def apply(self, fields):
+        """
+        Apply the individual limiters to specific prognostic variables
+        """
+
+        for _, suboption in self.suboptions.items():
+            field = fields.subfunctions[suboption.idx]
+            sublimiter.apply(field)
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
new file mode 100644
index 000000000..325ef645c
--- /dev/null
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -0,0 +1,313 @@
+"""
+This tests limiter options for different transport schemes.
+A sharp bubble of warm air is generated in a vertical slice and then transported
+by a prescribed transport scheme. If the limiter is working, the transport
+should have produced no new maxima or minima.
+"""
+
+from gusto import *
+from firedrake import (as_vector, PeriodicIntervalMesh, pi, SpatialCoordinate,
+                       ExtrudedMesh, FunctionSpace, Function, norm,
+                       conditional, sqrt)
+from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter
+import numpy as np
+import pytest
+
+
+def setup_limiters(dirname, space_A, space_B):
+
+    # ------------------------------------------------------------------------ #
+    # Parameters for test case
+    # ------------------------------------------------------------------------ #
+
+    Ld = 1.
+    tmax = 0.2
+    dt = tmax / 40
+    rotations = 0.25
+
+    # ------------------------------------------------------------------------ #
+    # Build model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    m = PeriodicIntervalMesh(20, Ld)
+    mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20))
+    degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] or space_B in ['DG0', 'Vtheta_degree_0'] else 1
+
+    domain = Domain(mesh, dt, family="CG", degree=degree)
+    
+    # Transporting velocity space
+    V = domain.spaces('HDiv')
+    
+    # Tracer A spaces
+    if space_A == 'DG0':
+        VA = domain.spaces('DG')
+        VCG1_A = FunctionSpace(mesh, 'CG', 1)
+        VDG1_A = domain.spaces('DG1_equispaced')
+        space_A_string = 'DG'
+    elif space_A == 'DG1':
+        VA = domain.spaces('DG')
+        space_A_string = 'DG'
+    elif space_A == 'DG1_equispaced':
+        VA = domain.spaces('DG1_equispaced')
+        space_A_string = 'DG1_equispaced'
+    elif space_A == 'Vtheta_degree_0':
+        VA = domain.spaces('theta')
+        VCG1_A = FunctionSpace(mesh, 'CG', 1)
+        VDG1_A = domain.spaces('DG1_equispaced')
+        space_A_string = 'theta'
+    elif space_A == 'Vtheta_degree_1':
+        VA = domain.spaces('theta')
+        space_A_string = 'theta'
+    else:
+        raise NotImplementedError
+    
+    # Tracer B spaces
+    if space_B == 'DG0':
+        VB = domain.spaces('DG')
+        VCG1_B = FunctionSpace(mesh, 'CG', 1)
+        VDG1_B = domain.spaces('DG1_equispaced')
+        space_B_string = 'DG'
+    elif space_B == 'DG1':
+        VB = domain.spaces('DG')
+        space_B_string = 'DG'
+    elif space_B == 'DG1_equispaced':
+        VB = domain.spaces('DG1_equispaced')
+        space_B_string = 'DG1_equispaced'
+    elif space_B == 'Vtheta_degree_0':
+        VB = domain.spaces('theta')
+        VCG1_B = FunctionSpace(mesh, 'CG', 1)
+        VDG1_B = domain.spaces('DG1_equispaced')
+        space_B_string = 'theta'
+    elif space_B == 'Vtheta_degree_1':
+        VB = domain.spaces('theta')
+        space_B_string = 'theta'
+    else:
+        raise NotImplementedError
+        
+    Vpsi = domain.spaces('H1')     
+        
+    tracerA = ActiveTracer(name='tracerA', space=space_A_string,
+                               variable_type=TracerVariableType.mixing_ratio,
+                               transport_eqn=TransportEquationType.advective)
+    
+    tracerB = ActiveTracer(name='tracerB', space=space_B_string,
+                               variable_type=TracerVariableType.mixing_ratio,
+                               transport_eqn=TransportEquationType.advective)
+        
+        
+    tracers = [tracerA, tracerB]
+
+    eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V)
+    output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1,
+                                  dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB'])
+    
+    io = IO(domain, output)
+
+    # ------------------------------------------------------------------------ #
+    # Set up transport scheme with options
+    # ------------------------------------------------------------------------ #
+
+    #suboptions = {}
+    sublimiters = {}
+
+    # Options and limiters for tracer_A
+
+    if space_A in ['DG0', 'Vtheta_degree_0']:
+        recover_opts = RecoveryOptions(embedding_space=VDG1_A,
+                               recovered_space=VCG1_A,
+                               project_low_method='recover',
+                               boundary_method=BoundaryMethod.taylor)
+                                   
+        sublimiters.update({'tracerA': VertexBasedLimiter(VDG1_A)})
+        
+    elif space_A == 'DG1':
+        sublimiters.update({'tracerA': DG1Limiter(VA)})
+
+    elif space_A == 'DG1_equispaced':
+        sublimiters.update({'tracerA': VertexBasedLimiter(VA)})
+
+    elif space_A == 'Vtheta_degree_1':
+        opts = EmbeddedDGOptions()
+        sublimiters.update({'tracerA': ThetaLimiter(VA)})
+
+    else:
+        raise NotImplementedError
+        
+    # Options and limiters for tracer_B
+
+    if space_B in ['DG0', 'Vtheta_degree_0']:
+        recover_opts = RecoveryOptions(embedding_space=VDG1_B,
+                               recovered_space=VCG1_B,
+                               project_low_method='recover',
+                               boundary_method=BoundaryMethod.taylor)
+                                   
+        sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
+        
+    elif space_B == 'DG1':
+        sublimiters.update({'tracerB': DG1Limiter(VB)})
+
+    elif space_B == 'DG1_equispaced':
+        sublimiters.update({'tracerB': VertexBasedLimiter(VB)})
+
+    elif space_B == 'Vtheta_degree_1':
+        opts = EmbeddedDGOptions()
+        sublimiters.update({'tracerB': ThetaLimiter(VB)})
+
+    else:
+        raise NotImplementedError
+
+    #MixedOptions = MixedFSOptions(eqn, suboptions)
+    MixedLimiter = MixedFSLimiter(eqn, sublimiters)
+
+    # Give the scheme for the coupled transport
+    transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
+    
+    # DG Upwind transport for both tracers:
+    transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
+    
+    # Build time stepper
+    stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method)
+
+    # ------------------------------------------------------------------------ #
+    # Initial condition
+    # ------------------------------------------------------------------------ #
+
+    tracerA_0 = stepper.fields('tracerA', space=VA)
+    tracerB_0 = stepper.fields('tracerB', space=VB)
+    true_fieldA = stepper.fields('true_tracerA', space=VA)
+    true_fieldB = stepper.fields('true_tracerB', space=VB)
+    
+    x, z = SpatialCoordinate(mesh)
+
+    tracer_min = 12.6
+    dtracer = 3.2
+
+    # First time do initial conditions, second time do final conditions
+    for i in range(2):
+
+        if i == 0:
+            x1_lower = 2 * Ld / 5
+            x1_upper = 3 * Ld / 5
+            z1_lower = 6 * Ld / 10
+            z1_upper = 8 * Ld / 10
+            x2_lower = 6 * Ld / 10
+            x2_upper = 8 * Ld / 10
+            z2_lower = 2 * Ld / 5
+            z2_upper = 3 * Ld / 5
+        elif i == 1:
+            # Rotated anti-clockwise by 90 degrees (x -> z, z -> -x)
+            x1_lower = 2 * Ld / 10
+            x1_upper = 4 * Ld / 10
+            z1_lower = 2 * Ld / 5
+            z1_upper = 3 * Ld / 5
+            x2_lower = 2 * Ld / 5
+            x2_upper = 3 * Ld / 5
+            z2_lower = 6 * Ld / 10
+            z2_upper = 8 * Ld / 10
+        else:
+            raise ValueError
+
+        expr_1 = conditional(x > x1_lower,
+                             conditional(x < x1_upper,
+                                         conditional(z > z1_lower,
+                                                     conditional(z < z1_upper, dtracer, 0.0),
+                                                     0.0),
+                                         0.0),
+                             0.0)
+
+        expr_2 = conditional(x > x2_lower,
+                             conditional(x < x2_upper,
+                                         conditional(z > z2_lower,
+                                                     conditional(z < z2_upper, dtracer, 0.0),
+                                                     0.0),
+                                         0.0),
+                             0.0)
+
+        if i == 0:
+            tracerA_0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+            tracerB_0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+        elif i == 1:
+            true_fieldA.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+            true_fieldB.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+        else:
+            raise ValueError
+
+    # ------------------------------------------------------------------------ #
+    # Velocity profile
+    # ------------------------------------------------------------------------ #
+
+    psi = Function(Vpsi)
+    u = stepper.fields('u')
+
+    # set up solid body rotation for transport
+    # we do this slightly complicated stream function to make the velocity 0 at edges
+    # thus we avoid odd effects at boundaries
+    xc = Ld / 2
+    zc = Ld / 2
+    r = sqrt((x - xc) ** 2 + (z - zc) ** 2)
+    omega = rotations * 2 * pi / tmax
+    r_out = 9 * Ld / 20
+    r_in = 2 * Ld / 5
+    A = omega * r_in / (2 * (r_in - r_out))
+    B = - omega * r_in * r_out / (r_in - r_out)
+    C = omega * r_in ** 2 * r_out / (r_in - r_out) / 2
+    psi_expr = conditional(r < r_in,
+                           omega * r ** 2 / 2,
+                           conditional(r < r_out,
+                                       A * r ** 2 + B * r + C,
+                                       A * r_out ** 2 + B * r_out + C))
+    psi.interpolate(psi_expr)
+
+    gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)])
+    u.project(gradperp(psi))
+
+    return stepper, tmax, true_fieldA, true_fieldB
+
+
+@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
+                                   'DG1', 'DG1_equispaced'])
+#Remove Dg1-dg1 and other easy ones after debugging
+@pytest.mark.parametrize('space_B', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
+                                    'DG1', 'DG1_equispaced'])
+
+
+def test_limiters(tmpdir, space_A, space_B):
+
+    # Setup and run
+    dirname = str(tmpdir)
+
+    stepper, tmax, true_fieldA, true_fieldB = setup_limiters(dirname, space_A, space_B)
+
+    stepper.run(t=0, tmax=tmax)
+
+    tol = 1e-9
+
+    final_fieldA = stepper.fields('tracerA')
+    final_fieldB = stepper.fields('tracerB')
+
+    # Check tracer is roughly in the correct place
+    assert norm(true_fieldA - final_fieldA) / norm(true_fieldA) < 0.05, \
+        'Something is wrong with the DG space tracer using a mixed limiter'
+
+    # Check tracer is roughly in the correct place
+    assert norm(true_fieldB - final_fieldB) / norm(true_fieldB) < 0.05, \
+        'Something is wrong with the DG1 equispaced tracer using a mixed limiter'
+
+    # Check for no new overshoots in A
+    assert np.max(final_fieldA.dat.data) <= np.max(true_fieldA.dat.data) + tol, \
+        'Application of the DG space limiter in the mixed limiter has not prevented overshoots'
+
+    # Check for no new undershoots in A
+    assert np.min(final_fieldA.dat.data) >= np.min(true_fieldA.dat.data) - tol, \
+        'Application of the DG space limiter in the mixed limiter has not prevented undershoots'
+
+    # Check for no new overshoots in B
+    assert np.max(final_fieldB.dat.data) <= np.max(true_fieldB.dat.data) + tol, \
+        'Application of the DG1 equispaced limiter in the mixed limiter has not prevented overshoots'
+
+    # Check for no new undershoots in B
+    assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \
+        'Application of the DG1 equispaced limiter in the mixed limiter has not prevented undershoots'
+
+    
\ No newline at end of file

From e04c73ac8ee9dcd5ee6bde0e1fe88f0636924b69 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Thu, 21 Dec 2023 20:18:24 +0000
Subject: [PATCH 03/38] MixedOptions dictionary added

---
 gusto/time_discretisation.py                  | 24 +++++------
 gusto/wrappers.py                             | 43 +++++++++++++------
 .../transport/test_mixed_fs_options.py        | 18 ++++----
 3 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 869b61931..ab3432b39 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -87,19 +87,19 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         self.courant_max = None
 
         if options is not None:
-            if options.wrapper_type = Mixed:
-                pass
-        else:
-            self.wrapper_name = options.name
-            if self.wrapper_name == "embedded_dg":
-                self.wrapper = EmbeddedDGWrapper(self, options)
-            elif self.wrapper_name == "recovered":
-                self.wrapper = RecoveryWrapper(self, options)
-            elif self.wrapper_name == "supg":
-                self.wrapper = SUPGWrapper(self, options)
+            if options.wrapper_type == 'mixed':
+                self.wrapper = options
             else:
-                raise RuntimeError(
-                    f'Time discretisation: wrapper {self.wrapper_name} not implemented')
+                self.wrapper_name = options.name
+                if self.wrapper_name == "embedded_dg":
+                    self.wrapper = EmbeddedDGWrapper(self, options)
+                elif self.wrapper_name == "recovered":
+                    self.wrapper = RecoveryWrapper(self, options)
+                elif self.wrapper_name == "supg":
+                    self.wrapper = SUPGWrapper(self, options)
+                else:
+                    raise RuntimeError(
+                        f'Time discretisation: wrapper {self.wrapper_name} not implemented')
         else:
             self.wrapper = None
             self.wrapper_name = None
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 1fd695685..71d97493e 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -381,19 +381,38 @@ def __init__(self, equation, suboptions):
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
     
-    self.suboptions = suboptions
-    self.wrapper_type = Mixed
-    
-    for field, suboption in suboptions:
-        # Check that the field is in the prognostic variable set:
-            if field not in equation.field_names:
-                raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
-            else:
-                # check that a valid wrapper has been given
-                
-                
-                self.suboptions[field].idx = equation.field_names.index(field)
+        self.suboptions = suboptions
+        self.wrapper_type = 'mixed'
+        self.wrappers = []
+        
+        for field, suboption in suboptions:
+            # Check that the field is in the prognostic variable set:
+                if field not in equation.field_names:
+                    raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
+                else:
+                    # check that a valid wrapper has been given
+                    wrapper_name = suboption.name
+                    
+                    if wrapper_name == "embedded_dg":
+                        self.wrappers.append(EmbeddedDGWrapper(self, options))
+                        self.wrapper_fields.append(field)
+                    elif wrapper_name == "recovered":
+                        self.wrappers.append(RecoveryWrapper(self, options))
+                        self.wrapper_fields.append(field)
+                    elif wrapper_name == "supg":
+                        self.wrappers.append(SUPGWrapper(self, options))
+                        self.wrapper_fields.append(field)
+                    else:
+                        raise RuntimeError(
+                        f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
+                    
+                    #Initialise the wrapper and associate with a field:
+                    
+                    self.suboptions[field].idx = equation.field_names.index(field)
     
+    def setup(self):
+        # This is done in the suboption wrappers themselves
+        pass
     
     def apply(self, fields):
         """
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 325ef645c..94ec51bb6 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -108,16 +108,16 @@ def setup_limiters(dirname, space_A, space_B):
     # Set up transport scheme with options
     # ------------------------------------------------------------------------ #
 
-    #suboptions = {}
+    suboptions = {}
     sublimiters = {}
 
     # Options and limiters for tracer_A
 
     if space_A in ['DG0', 'Vtheta_degree_0']:
-        recover_opts = RecoveryOptions(embedding_space=VDG1_A,
-                               recovered_space=VCG1_A,
-                               project_low_method='recover',
-                               boundary_method=BoundaryMethod.taylor)
+        suboptions.update({'tracerA': RecoveryOptions(embedding_space=VDG1_A,
+                                              recovered_space=VCG1_A,
+                                              project_low_method='recover',
+                                              boundary_method=BoundaryMethod.taylor)})
                                    
         sublimiters.update({'tracerA': VertexBasedLimiter(VDG1_A)})
         
@@ -157,11 +157,11 @@ def setup_limiters(dirname, space_A, space_B):
     else:
         raise NotImplementedError
 
-    #MixedOptions = MixedFSOptions(eqn, suboptions)
+    opts = MixedOptions(eqn, suboptions)
     MixedLimiter = MixedFSLimiter(eqn, sublimiters)
 
     # Give the scheme for the coupled transport
-    transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
+    transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter)
     
     # DG Upwind transport for both tracers:
     transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
@@ -268,8 +268,8 @@ def setup_limiters(dirname, space_A, space_B):
 @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
                                    'DG1', 'DG1_equispaced'])
 #Remove Dg1-dg1 and other easy ones after debugging
-@pytest.mark.parametrize('space_B', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-                                    'DG1', 'DG1_equispaced'])
+@pytest.mark.parametrize('space_B', ['Vtheta_degree_0'])#, 'Vtheta_degree_1', 'DG0',
+                                    #'DG1', 'DG1_equispaced'])
 
 
 def test_limiters(tmpdir, space_A, space_B):

From 4a13df3a326ce8404ec18ef655a523e289885280 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Fri, 22 Dec 2023 02:43:21 +0000
Subject: [PATCH 04/38] improvements to MixedOptions

---
 gusto/time_discretisation.py                  |  56 ++++++---
 gusto/wrappers.py                             |  47 +++++---
 .../transport/test_mixed_fs_options.py        | 106 ++++++++++--------
 3 files changed, 131 insertions(+), 78 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index ab3432b39..593bbe830 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -37,6 +37,20 @@ def wrapper_apply(original_apply):
     """Decorator to add steps for using a wrapper around the apply method."""
     def get_apply(self, x_out, x_in):
 
+        #if type(self.wrapper) == MixedOptions:
+        
+            # Subwrappers for different tracers have been defined
+            # Use the apply defined in the MixedOptions object
+            
+         #   print('aha')
+         #   self.wrapper.pre_apply(x_in)
+            
+         #   for _, subwrapper in self.wrapper.suboptions.items(): 
+         #       original_apply(self, self.subwrapper.x_out, self.subwrapper.x_in)
+            
+         #   self.wrapper.post_apply(x_out)
+            
+
         if self.wrapper is not None:
 
             def new_apply(self, x_out, x_in):
@@ -87,8 +101,11 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         self.courant_max = None
 
         if options is not None:
-            if options.wrapper_type == 'mixed':
+            if type(options) == MixedOptions:
+                print('jahjah')
                 self.wrapper = options
+                # Or do I need to initialise everything here
+                # like is done for a single wrapper?
             else:
                 self.wrapper_name = options.name
                 if self.wrapper_name == "embedded_dg":
@@ -162,21 +179,28 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         # -------------------------------------------------------------------- #
 
         if self.wrapper is not None:
-            self.wrapper.setup()
-            self.fs = self.wrapper.function_space
-            if self.solver_parameters is None:
-                self.solver_parameters = self.wrapper.solver_parameters
-            new_test = TestFunction(self.wrapper.test_space)
-            # SUPG has a special wrapper
-            if self.wrapper_name == "supg":
-                new_test = self.wrapper.test
-
-            # Replace the original test function with the one from the wrapper
-            self.residual = self.residual.label_map(
-                all_terms,
-                map_if_true=replace_test_function(new_test))
-
-            self.residual = self.wrapper.label_terms(self.residual)
+            if type(self.wrapper) == MixedOptions:
+                # Subwrappers are defined.
+                # Set these up with ?
+                for _, subwrapper in self.wrapper.suboptions.items():
+                    print(subwrapper)
+                pass
+            else:
+                self.wrapper.setup()
+                self.fs = self.wrapper.function_space
+                if self.solver_parameters is None:
+                    self.solver_parameters = self.wrapper.solver_parameters
+                new_test = TestFunction(self.wrapper.test_space)
+                # SUPG has a special wrapper
+                if self.wrapper_name == "supg":
+                    new_test = self.wrapper.test
+    
+                # Replace the original test function with the one from the wrapper
+                self.residual = self.residual.label_map(
+                    all_terms,
+                    map_if_true=replace_test_function(new_test))
+    
+                self.residual = self.wrapper.label_terms(self.residual)
 
         # -------------------------------------------------------------------- #
         # Make boundary conditions
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 71d97493e..b14a23b6b 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -380,28 +380,33 @@ def __init__(self, equation, suboptions):
         Raises:
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
-    
+        print(equation.function_space)
+        self.x_in = Function(equation.function_space)
+        self.x_out = Function(equation.function_space)
+        
+        print('Initialising MixedOptions')
+        
         self.suboptions = suboptions
-        self.wrapper_type = 'mixed'
-        self.wrappers = []
         
-        for field, suboption in suboptions:
+        print(suboptions.items())
+        
+        self.wrapper = {}
+        
+        for field, suboption in suboptions.items():
             # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
                 else:
                     # check that a valid wrapper has been given
                     wrapper_name = suboption.name
+                    print(wrapper_name)
                     
                     if wrapper_name == "embedded_dg":
-                        self.wrappers.append(EmbeddedDGWrapper(self, options))
-                        self.wrapper_fields.append(field)
+                        self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)})
                     elif wrapper_name == "recovered":
-                        self.wrappers.append(RecoveryWrapper(self, options))
-                        self.wrapper_fields.append(field)
+                        self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption)
                     elif wrapper_name == "supg":
-                        self.wrappers.append(SUPGWrapper(self, options))
-                        self.wrapper_fields.append(field)
+                        self.suboptions[field].subwrapper = SUPGWrapper(self, suboption)
                     else:
                         raise RuntimeError(
                         f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
@@ -409,16 +414,30 @@ def __init__(self, equation, suboptions):
                     #Initialise the wrapper and associate with a field:
                     
                     self.suboptions[field].idx = equation.field_names.index(field)
+                    #self.suboptions[field].fs = equation.field_names.function_space(field)
+ #                   self.subwrapper.x_in = Function(self.suboptions[field].fs)
     
     def setup(self):
         # This is done in the suboption wrappers themselves
         pass
     
-    def apply(self, fields):
+    def pre_apply(self, x_in):
         """
-        Apply the individual limiters to specific prognostic variables
+        Perform the pre-applications from all subwrappers
+        """
+
+        for _, subwrapper in self.suboptions.items():
+            print(subwrapper)
+            field = x_in.subfunctions[subwrapper.idx]
+            subwrapper.pre_apply(field)
+            
+    def post_apply(self, x_in):
+        """
+        Perform the post-applications from all subwrappers
         """
 
         for _, suboption in self.suboptions.items():
-            field = fields.subfunctions[suboption.idx]
-            sublimiter.apply(field)
+            print(suboption)
+            field = x_in.subfunctions[suboption.idx]
+            suboption.subwrapper.post_apply(field)
+
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 94ec51bb6..a82136923 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -32,7 +32,7 @@ def setup_limiters(dirname, space_A, space_B):
     # Domain
     m = PeriodicIntervalMesh(20, Ld)
     mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20))
-    degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] or space_B in ['DG0', 'Vtheta_degree_0'] else 1
+    degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] else 1
 
     domain = Domain(mesh, dt, family="CG", degree=degree)
     
@@ -63,25 +63,28 @@ def setup_limiters(dirname, space_A, space_B):
         raise NotImplementedError
     
     # Tracer B spaces
-    if space_B == 'DG0':
-        VB = domain.spaces('DG')
-        VCG1_B = FunctionSpace(mesh, 'CG', 1)
-        VDG1_B = domain.spaces('DG1_equispaced')
-        space_B_string = 'DG'
-    elif space_B == 'DG1':
-        VB = domain.spaces('DG')
-        space_B_string = 'DG'
-    elif space_B == 'DG1_equispaced':
-        VB = domain.spaces('DG1_equispaced')
-        space_B_string = 'DG1_equispaced'
-    elif space_B == 'Vtheta_degree_0':
-        VB = domain.spaces('theta')
-        VCG1_B = FunctionSpace(mesh, 'CG', 1)
-        VDG1_B = domain.spaces('DG1_equispaced')
-        space_B_string = 'theta'
-    elif space_B == 'Vtheta_degree_1':
-        VB = domain.spaces('theta')
-        space_B_string = 'theta'
+    if space_B == 'DG':
+        if degree == 0:
+            VB = domain.spaces('DG')
+            VCG1_B = FunctionSpace(mesh, 'CG', 1)
+            VDG1_B = domain.spaces('DG1_equispaced')
+            space_B_string = 'DG'
+        elif degree == 1:
+            VB = domain.spaces('DG')
+            space_B_string = 'DG'
+        else:
+            raise NotImplementedError
+    elif space_B == 'Vtheta':
+        if degree == 0: 
+            VB = domain.spaces('theta')
+            VCG1_B = FunctionSpace(mesh, 'CG', 1)
+            VDG1_B = domain.spaces('DG1_equispaced')
+            space_B_string = 'theta'
+        elif degree == 1:
+            VB = domain.spaces('theta')
+            space_B_string = 'theta'
+        else:
+            raise NotImplementedError
     else:
         raise NotImplementedError
         
@@ -128,32 +131,39 @@ def setup_limiters(dirname, space_A, space_B):
         sublimiters.update({'tracerA': VertexBasedLimiter(VA)})
 
     elif space_A == 'Vtheta_degree_1':
-        opts = EmbeddedDGOptions()
+        suboptions.update({'tracerA': EmbeddedDGOptions()})
         sublimiters.update({'tracerA': ThetaLimiter(VA)})
 
     else:
         raise NotImplementedError
         
     # Options and limiters for tracer_B
-
-    if space_B in ['DG0', 'Vtheta_degree_0']:
-        recover_opts = RecoveryOptions(embedding_space=VDG1_B,
-                               recovered_space=VCG1_B,
-                               project_low_method='recover',
-                               boundary_method=BoundaryMethod.taylor)
+    
+    if space_B == 'DG':
+        if degree == 0:
+            suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B,
+                                              recovered_space=VCG1_B,
+                                              project_low_method='recover',
+                                              boundary_method=BoundaryMethod.taylor)})
                                    
-        sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
-        
-    elif space_B == 'DG1':
-        sublimiters.update({'tracerB': DG1Limiter(VB)})
-
-    elif space_B == 'DG1_equispaced':
-        sublimiters.update({'tracerB': VertexBasedLimiter(VB)})
-
-    elif space_B == 'Vtheta_degree_1':
-        opts = EmbeddedDGOptions()
-        sublimiters.update({'tracerB': ThetaLimiter(VB)})
-
+            sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
+        elif degree == 1:
+            sublimiters.update({'tracerB': DG1Limiter(VB)})
+        else:
+            raise NotImplementedError
+    elif space_B == 'Vtheta':
+        if degree == 0: 
+            suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B,
+                                              recovered_space=VCG1_B,
+                                              project_low_method='recover',
+                                              boundary_method=BoundaryMethod.taylor)})
+                                   
+            sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
+        elif degree == 1:
+            opts = EmbeddedDGOptions()
+            sublimiters.update({'tracerB': ThetaLimiter(VB)})
+        else:
+            raise NotImplementedError
     else:
         raise NotImplementedError
 
@@ -162,6 +172,7 @@ def setup_limiters(dirname, space_A, space_B):
 
     # Give the scheme for the coupled transport
     transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter)
+    #transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
     
     # DG Upwind transport for both tracers:
     transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
@@ -267,9 +278,8 @@ def setup_limiters(dirname, space_A, space_B):
 
 @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
                                    'DG1', 'DG1_equispaced'])
-#Remove Dg1-dg1 and other easy ones after debugging
-@pytest.mark.parametrize('space_B', ['Vtheta_degree_0'])#, 'Vtheta_degree_1', 'DG0',
-                                    #'DG1', 'DG1_equispaced'])
+# It only makes sense to use the same degree for tracer B
+@pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 
 
 def test_limiters(tmpdir, space_A, space_B):
@@ -288,26 +298,26 @@ def test_limiters(tmpdir, space_A, space_B):
 
     # Check tracer is roughly in the correct place
     assert norm(true_fieldA - final_fieldA) / norm(true_fieldA) < 0.05, \
-        'Something is wrong with the DG space tracer using a mixed limiter'
+        f"Something is wrong with the {space_A} space tracer using a mixed limiter"
 
     # Check tracer is roughly in the correct place
     assert norm(true_fieldB - final_fieldB) / norm(true_fieldB) < 0.05, \
-        'Something is wrong with the DG1 equispaced tracer using a mixed limiter'
+        f"Something is wrong with the {space_B} space tracer using a mixed limiter"
 
     # Check for no new overshoots in A
     assert np.max(final_fieldA.dat.data) <= np.max(true_fieldA.dat.data) + tol, \
-        'Application of the DG space limiter in the mixed limiter has not prevented overshoots'
+        f"Application of the {space_A} space limiter in the mixed limiter has not prevented overshoots"
 
     # Check for no new undershoots in A
     assert np.min(final_fieldA.dat.data) >= np.min(true_fieldA.dat.data) - tol, \
-        'Application of the DG space limiter in the mixed limiter has not prevented undershoots'
+        f"Application of the {space_B} space limiter in the mixed limiter has not prevented undershoots"
 
     # Check for no new overshoots in B
     assert np.max(final_fieldB.dat.data) <= np.max(true_fieldB.dat.data) + tol, \
-        'Application of the DG1 equispaced limiter in the mixed limiter has not prevented overshoots'
+        f"Application of the {space_A} limiter in the mixed limiter has not prevented overshoots"
 
     # Check for no new undershoots in B
     assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \
-        'Application of the DG1 equispaced limiter in the mixed limiter has not prevented undershoots'
+        f"Application of the {space_B} equispaced limiter in the mixed limiter has not prevented undershoots"
 
     
\ No newline at end of file

From 616505790ecbe48cdbe4ebfa7478fca8598d21f0 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 9 Jan 2024 04:01:49 +0000
Subject: [PATCH 05/38] more mixed options changes

---
 gusto/time_discretisation.py                  | 44 +++++++++++-
 gusto/wrappers.py                             | 69 ++++++++++++-------
 .../transport/test_mixed_fs_options.py        |  7 +-
 3 files changed, 91 insertions(+), 29 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 593bbe830..c537b317f 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -56,7 +56,10 @@ def get_apply(self, x_out, x_in):
             def new_apply(self, x_out, x_in):
 
                 self.wrapper.pre_apply(x_in)
-                original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
+                if type(self.wrapper) == MixedOptions:
+                    original_apply(self, x_out, x_in)
+                else:
+                    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
                 self.wrapper.post_apply(x_out)
 
             return new_apply(self, x_out, x_in)
@@ -103,9 +106,39 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         if options is not None:
             if type(options) == MixedOptions:
                 print('jahjah')
+                print(options)
                 self.wrapper = options
+                #self.subwrappers = {}
                 # Or do I need to initialise everything here
                 # like is done for a single wrapper?
+                for field, suboption in self.wrapper.suboptions.items():
+                    print(field)
+                    print(suboption)
+                    
+                    #Replace options with a wrapper?
+                    
+                    #if suboption.name == 'embedded_dg':
+                    #    self.suboptions[field].wrapper = EmbeddedDGWrapper(self, suboption)
+                    #elif suboption.name == "recovered":
+                    #    self.suboptions[field].wrapper = RecoveryWrapper(self, suboption)
+                    #elif suboption.name == "supg":
+                    #    self.suboptions[field].wrapper = SUPGWrapper(self, suboption)
+                    #else:
+                    #    raise RuntimeError(
+                    #    f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
+                    if suboption.name == 'embedded_dg':
+                        self.wrapper.subwrappers.update({field:EmbeddedDGWrapper(self, suboption)})
+                    elif suboption.name == "recovered":
+                        self.wrapper.subwrappers.update({field:RecoveryWrapper(self, suboption)})
+                    elif suboption.name == "supg":
+                        self.wrapper.subwrappers.update({field:SUPGWrapper(self, suboption)})
+                    else:
+                        raise RuntimeError(
+                        f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
+                print(self.wrapper)
+                print(self.wrapper.suboptions)
+                print(self.wrapper.subwrappers)
+                #self.wrapper_name = 'mixed'
             else:
                 self.wrapper_name = options.name
                 if self.wrapper_name == "embedded_dg":
@@ -180,10 +213,17 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         if self.wrapper is not None:
             if type(self.wrapper) == MixedOptions:
+            #if self.wrapper_name == 'mixed':
                 # Subwrappers are defined.
                 # Set these up with ?
-                for _, subwrapper in self.wrapper.suboptions.items():
+                for field, subwrapper in self.wrapper.subwrappers.items():
+                    #Set up field idxs here.
+                    print(field)
                     print(subwrapper)
+                    self.wrapper.subwrappers[field].idx = equation.field_names.index(field)
+                    self.wrapper.subwrappers[field].mixed_options = True
+                    self.wrapper.subwrappers[field].setup()
+                    self.fs = self.wrapper.subwrappers[field].function_space
                 pass
             else:
                 self.wrapper.setup()
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index b14a23b6b..52da12e4f 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -33,7 +33,9 @@ def __init__(self, time_discretisation, wrapper_options):
         self.time_discretisation = time_discretisation
         self.options = wrapper_options
         self.solver_parameters = None
-        self.field_name = None
+        #self.field_name = None
+        self.idx = None
+        self.mixed_options = False
 
     @abstractmethod
     def setup(self):
@@ -162,8 +164,10 @@ class RecoveryWrapper(Wrapper):
     def setup(self):
         """Sets up function spaces and fields needed for this wrapper."""
 
+        print(self.options)
+
         assert isinstance(self.options, RecoveryOptions), \
-            'Embedded DG wrapper can only be used with Recovery Options'
+            'Recovery wrapper can only be used with Recovery Options'
 
         original_space = self.time_discretisation.fs
         domain = self.time_discretisation.domain
@@ -177,6 +181,7 @@ def setup(self):
             V_elt = BrokenElement(original_space.ufl_element())
             self.function_space = FunctionSpace(domain.mesh, V_elt)
         else:
+            print('using embedded space')
             self.function_space = self.options.embedding_space
 
         self.test_space = self.function_space
@@ -185,11 +190,23 @@ def setup(self):
         # Internal variables to be used
         # -------------------------------------------------------------------- #
 
-        self.x_in_tmp = Function(self.time_discretisation.fs)
+        if self.mixed_options == True:
+            self.x_in_tmp = Function(self.function_space)
+        else:
+            self.x_in_tmp = Function(self.time_discretisation.fs)
+            
+        print(self.function_space)
+        print(self.time_discretisation.fs)
+            
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
+        
+        #self.x_out = Function(equation.function_space)
+        #self.x_out = Function(self.time_discretisation.fs)
+        
         if self.time_discretisation.idx is None:
-            self.x_projected = Function(equation.function_space)
+            #self.x_projected = Function(equation.function_space)
+            self.x_projected = Function(self.function_space)
         else:
             self.x_projected = Function(equation.spaces[self.time_discretisation.idx])
 
@@ -381,18 +398,20 @@ def __init__(self, equation, suboptions):
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
         print(equation.function_space)
-        self.x_in = Function(equation.function_space)
-        self.x_out = Function(equation.function_space)
+        #self.x_in = Function(equation.function_space)
+        #self.x_out = Function(equation.function_space)
         
         print('Initialising MixedOptions')
         
         self.suboptions = suboptions
         
-        print(suboptions.items())
+        #print(suboptions.items())
         
-        self.wrapper = {}
+        self.subwrappers = {}
         
         for field, suboption in suboptions.items():
+                print(field)
+                print(suboption)
             # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
@@ -401,19 +420,19 @@ def __init__(self, equation, suboptions):
                     wrapper_name = suboption.name
                     print(wrapper_name)
                     
-                    if wrapper_name == "embedded_dg":
-                        self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)})
-                    elif wrapper_name == "recovered":
-                        self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption)
-                    elif wrapper_name == "supg":
-                        self.suboptions[field].subwrapper = SUPGWrapper(self, suboption)
-                    else:
-                        raise RuntimeError(
-                        f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
+                    #if wrapper_name == "embedded_dg":
+                    #    self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)})
+                    #elif wrapper_name == "recovered":
+                    #    self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption)
+                    #elif wrapper_name == "supg":
+                    #    self.suboptions[field].subwrapper = SUPGWrapper(self, suboption)
+                    #else:
+                    #    raise RuntimeError(
+                    #    f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
                     
                     #Initialise the wrapper and associate with a field:
-                    
-                    self.suboptions[field].idx = equation.field_names.index(field)
+                    #self.suboptions[field].wrapper_name = suboption.name
+                    #self.suboptions[field].idx = equation.field_names.index(field)
                     #self.suboptions[field].fs = equation.field_names.function_space(field)
  #                   self.subwrapper.x_in = Function(self.suboptions[field].fs)
     
@@ -426,8 +445,9 @@ def pre_apply(self, x_in):
         Perform the pre-applications from all subwrappers
         """
 
-        for _, subwrapper in self.suboptions.items():
+        for _, subwrapper in self.subwrappers.items():
             print(subwrapper)
+            print('pre')
             field = x_in.subfunctions[subwrapper.idx]
             subwrapper.pre_apply(field)
             
@@ -436,8 +456,9 @@ def post_apply(self, x_in):
         Perform the post-applications from all subwrappers
         """
 
-        for _, suboption in self.suboptions.items():
-            print(suboption)
-            field = x_in.subfunctions[suboption.idx]
-            suboption.subwrapper.post_apply(field)
+        for _, subwrapper in self.subwrappers.items():
+            print(subwrapper)
+            print('post')
+            field = x_in.subfunctions[subwrapper.idx]
+            subwrapper.post_apply(field)
 
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index a82136923..9701e3dac 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -160,7 +160,7 @@ def setup_limiters(dirname, space_A, space_B):
                                    
             sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
         elif degree == 1:
-            opts = EmbeddedDGOptions()
+            suboptions.update({'tracerB': EmbeddedDGOptions()})
             sublimiters.update({'tracerB': ThetaLimiter(VB)})
         else:
             raise NotImplementedError
@@ -276,8 +276,9 @@ def setup_limiters(dirname, space_A, space_B):
     return stepper, tmax, true_fieldA, true_fieldB
 
 
-@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-                                   'DG1', 'DG1_equispaced'])
+#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
+#                                   'DG1', 'DG1_equispaced'])
+@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced'])
 # It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 

From 775d3b2101c73f4b3a3de040de73a587d5010b10 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 10 Jan 2024 01:39:48 +0000
Subject: [PATCH 06/38] more changes to mixed options

---
 gusto/equations.py                            |  3 +
 gusto/limiters.py                             |  2 +
 gusto/time_discretisation.py                  | 33 ++++++++--
 gusto/wrappers.py                             | 66 ++++++++++++++-----
 .../transport/test_mixed_fs_options.py        |  9 +--
 5 files changed, 87 insertions(+), 26 deletions(-)

diff --git a/gusto/equations.py b/gusto/equations.py
index bd61e795b..c30124c73 100644
--- a/gusto/equations.py
+++ b/gusto/equations.py
@@ -426,6 +426,8 @@ def add_tracers_to_prognostics(self, domain, active_tracers):
                 else:
                     self.space_names[tracer.name] = tracer.space
                 self.spaces.append(domain.spaces(tracer.space))
+                #print('self.spaces is', self.spaces)
+                print(len(self.spaces))
             else:
                 raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}')
 
@@ -542,6 +544,7 @@ def __init__(self, domain, active_tracers, Vu=None):
         # Build finite element spaces
         self.spaces = []
 
+
         # Add active tracers to the list of prognostics
         if active_tracers is None:
             active_tracers = []
diff --git a/gusto/limiters.py b/gusto/limiters.py
index a39c6688d..16a2d3b40 100644
--- a/gusto/limiters.py
+++ b/gusto/limiters.py
@@ -203,6 +203,8 @@ def __init__(self, equation, sublimiters):
                 raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
             else:
                 self.sublimiters[field].idx = equation.field_names.index(field)
+                
+                # Give the sublimiter a subspace?
 
     def apply(self, fields):
         """
diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index c537b317f..ad45fd84c 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -57,7 +57,8 @@ def new_apply(self, x_out, x_in):
 
                 self.wrapper.pre_apply(x_in)
                 if type(self.wrapper) == MixedOptions:
-                    original_apply(self, x_out, x_in)
+                    #original_apply(self, x_out, x_in)
+                    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
                 else:
                     original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
                 self.wrapper.post_apply(x_out)
@@ -105,8 +106,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
 
         if options is not None:
             if type(options) == MixedOptions:
-                print('jahjah')
-                print(options)
+                #print('jahjah')
+                #print(options)
                 self.wrapper = options
                 #self.subwrappers = {}
                 # Or do I need to initialise everything here
@@ -216,14 +217,36 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             #if self.wrapper_name == 'mixed':
                 # Subwrappers are defined.
                 # Set these up with ?
+                
+                # Give more than one fs?
+                fields = []
+                
+                print(self.wrapper.wrapper_spaces)
+                
                 for field, subwrapper in self.wrapper.subwrappers.items():
-                    #Set up field idxs here.
+                    # Set up field idxs here.
                     print(field)
                     print(subwrapper)
                     self.wrapper.subwrappers[field].idx = equation.field_names.index(field)
                     self.wrapper.subwrappers[field].mixed_options = True
+                    
+                    # Store the original space of the tracer
+                    self.wrapper.subwrappers[field].tracer_fs = self.equation.spaces[equation.field_names.index(field)]
                     self.wrapper.subwrappers[field].setup()
-                    self.fs = self.wrapper.subwrappers[field].function_space
+                    
+                    # Update the function space to that needed by the wrapper
+                    self.wrapper.wrapper_spaces.update({field: self.wrapper.subwrappers[field].function_space})
+                    fields.append(self.wrapper.subwrappers[field].function_space)
+                    # Append the new function space from the wrapper
+                    #print('self.tracer_fs is', self.wrapper.subwrappers[field].tracer_fs)
+                print(self.wrapper.wrapper_spaces)
+                #Setup the mixed wrapper
+                #yo_spaces = []
+                #for _, space_name in self.wrapper.wrapper_spaces.items():
+                #    yo_spaces.append(self.domain.spaces(space_name))
+                #print(yo_spaces)
+                self.wrapper.wrapper_spaces = fields
+                self.wrapper.setup()
                 pass
             else:
                 self.wrapper.setup()
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 52da12e4f..d4fa6d4dd 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -7,7 +7,7 @@
 from abc import ABCMeta, abstractmethod
 from firedrake import (
     FunctionSpace, Function, BrokenElement, Projector, Interpolator,
-    VectorElement, Constant, as_ufl, dot, grad, TestFunction
+    VectorElement, Constant, as_ufl, dot, grad, TestFunction, MixedFunctionSpace
 )
 from firedrake.fml import Term
 from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions
@@ -36,6 +36,7 @@ def __init__(self, time_discretisation, wrapper_options):
         #self.field_name = None
         self.idx = None
         self.mixed_options = False
+        self.tracer_fs = None
 
     @abstractmethod
     def setup(self):
@@ -84,10 +85,16 @@ def setup(self):
 
         assert isinstance(self.options, EmbeddedDGOptions), \
             'Embedded DG wrapper can only be used with Embedded DG Options'
-
-        original_space = self.time_discretisation.fs
+        
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
+        
+        if self.mixed_options == True:
+            #print('self.tracer_fs is', self.tracer_fs)
+            #print('self.time_discretisation.fs is', self.time_discretisation.fs)
+            original_space = self.tracer_fs
+        else:
+            original_space = self.time_discretisation.fs
 
         # -------------------------------------------------------------------- #
         # Set up spaces to be used with wrapper
@@ -107,7 +114,9 @@ def setup(self):
 
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
-        if self.time_discretisation.idx is None:
+        if self.mixed_options == True:
+            self.x_projected = Function(self.tracer_fs)
+        elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
         else:
             self.x_projected = Function(equation.spaces[self.time_discretisation.idx])
@@ -169,9 +178,13 @@ def setup(self):
         assert isinstance(self.options, RecoveryOptions), \
             'Recovery wrapper can only be used with Recovery Options'
 
-        original_space = self.time_discretisation.fs
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
+        
+        if self.mixed_options == True:
+            original_space = self.tracer_fs
+        else:
+            original_space = self.time_discretisation.fs
 
         # -------------------------------------------------------------------- #
         # Set up spaces to be used with wrapper
@@ -191,12 +204,13 @@ def setup(self):
         # -------------------------------------------------------------------- #
 
         if self.mixed_options == True:
-            self.x_in_tmp = Function(self.function_space)
+            self.x_in_tmp = Function(self.tracer_fs)
+            #self.x_in_tmp = Function(equation.spaces[self.idx])
         else:
             self.x_in_tmp = Function(self.time_discretisation.fs)
             
-        print(self.function_space)
-        print(self.time_discretisation.fs)
+        #print(self.function_space)
+        #print(self.time_discretisation.fs)
             
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
@@ -204,9 +218,10 @@ def setup(self):
         #self.x_out = Function(equation.function_space)
         #self.x_out = Function(self.time_discretisation.fs)
         
-        if self.time_discretisation.idx is None:
-            #self.x_projected = Function(equation.function_space)
-            self.x_projected = Function(self.function_space)
+        if self.mixed_options == True:
+            self.x_projected = Function(self.tracer_fs)
+        elif self.time_discretisation.idx is None:
+            self.x_projected = Function(equation.function_space)
         else:
             self.x_projected = Function(equation.spaces[self.time_discretisation.idx])
 
@@ -397,7 +412,14 @@ def __init__(self, equation, suboptions):
         Raises:
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
-        print(equation.function_space)
+        #print(equation.function_space)
+        #print(equation.active_tracers)
+        #print(equation.active_tracers[0])
+        #print(equation.active_tracers[1])
+        print(equation.space_names)
+        
+        self.wrapper_spaces = equation.space_names
+        
         #self.x_in = Function(equation.function_space)
         #self.x_out = Function(equation.function_space)
         
@@ -412,13 +434,17 @@ def __init__(self, equation, suboptions):
         for field, suboption in suboptions.items():
                 print(field)
                 print(suboption)
-            # Check that the field is in the prognostic variable set:
+                
+                # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
                 else:
                     # check that a valid wrapper has been given
                     wrapper_name = suboption.name
-                    print(wrapper_name)
+                    #print(wrapper_name)
+                    
+                    # Extract the space?
+                    
                     
                     #if wrapper_name == "embedded_dg":
                     #    self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)})
@@ -438,7 +464,13 @@ def __init__(self, equation, suboptions):
     
     def setup(self):
         # This is done in the suboption wrappers themselves
-        pass
+        # Or, determine the new mixed function space
+        self.function_space = MixedFunctionSpace(self.wrapper_spaces)
+        self.x_in = Function(self.function_space)
+        self.x_out = Function(self.function_space) 
+        
+        
+        #pass
     
     def pre_apply(self, x_in):
         """
@@ -450,6 +482,7 @@ def pre_apply(self, x_in):
             print('pre')
             field = x_in.subfunctions[subwrapper.idx]
             subwrapper.pre_apply(field)
+            print('success')
             
     def post_apply(self, x_in):
         """
@@ -460,5 +493,4 @@ def post_apply(self, x_in):
             print(subwrapper)
             print('post')
             field = x_in.subfunctions[subwrapper.idx]
-            subwrapper.post_apply(field)
-
+            subwrapper.post_apply(field)
\ No newline at end of file
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 9701e3dac..699618a64 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -171,6 +171,7 @@ def setup_limiters(dirname, space_A, space_B):
     MixedLimiter = MixedFSLimiter(eqn, sublimiters)
 
     # Give the scheme for the coupled transport
+    #transport_schemes = SSPRK3(domain, options=opts)
     transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter)
     #transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
     
@@ -276,14 +277,14 @@ def setup_limiters(dirname, space_A, space_B):
     return stepper, tmax, true_fieldA, true_fieldB
 
 
-#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-#                                   'DG1', 'DG1_equispaced'])
-@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced'])
+@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
+                                   'DG1', 'DG1_equispaced'])
+#@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced'])
 # It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 
 
-def test_limiters(tmpdir, space_A, space_B):
+def test_mixed_fs_options(tmpdir, space_A, space_B):
 
     # Setup and run
     dirname = str(tmpdir)

From 793768ef7919b8832dfb4cc5cd495982c3ae83ab Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 10 Jan 2024 22:58:18 +0000
Subject: [PATCH 07/38] working on new test functions with mixed options

---
 gusto/time_discretisation.py | 97 ++++++++++++++++++------------------
 gusto/wrappers.py            | 42 +++++++++++-----
 2 files changed, 78 insertions(+), 61 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index ad45fd84c..90f73a502 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -37,30 +37,21 @@ def wrapper_apply(original_apply):
     """Decorator to add steps for using a wrapper around the apply method."""
     def get_apply(self, x_out, x_in):
 
-        #if type(self.wrapper) == MixedOptions:
-        
-            # Subwrappers for different tracers have been defined
-            # Use the apply defined in the MixedOptions object
-            
-         #   print('aha')
-         #   self.wrapper.pre_apply(x_in)
-            
-         #   for _, subwrapper in self.wrapper.suboptions.items(): 
-         #       original_apply(self, self.subwrapper.x_out, self.subwrapper.x_in)
-            
-         #   self.wrapper.post_apply(x_out)
-            
-
         if self.wrapper is not None:
 
             def new_apply(self, x_out, x_in):
 
                 self.wrapper.pre_apply(x_in)
-                if type(self.wrapper) == MixedOptions:
+                original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
+                
+                #if type(self.wrapper) == MixedOptions:
+                    #print('x_out', x_out)
+                    #print('self.wrapper.x_out', self.wrapper.x_out)
                     #original_apply(self, x_out, x_in)
-                    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
-                else:
-                    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
+                #    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
+                #else:
+                #    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
+                
                 self.wrapper.post_apply(x_out)
 
             return new_apply(self, x_out, x_in)
@@ -116,17 +107,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                     print(field)
                     print(suboption)
                     
-                    #Replace options with a wrapper?
-                    
-                    #if suboption.name == 'embedded_dg':
-                    #    self.suboptions[field].wrapper = EmbeddedDGWrapper(self, suboption)
-                    #elif suboption.name == "recovered":
-                    #    self.suboptions[field].wrapper = RecoveryWrapper(self, suboption)
-                    #elif suboption.name == "supg":
-                    #    self.suboptions[field].wrapper = SUPGWrapper(self, suboption)
-                    #else:
-                    #    raise RuntimeError(
-                    #    f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
                     if suboption.name == 'embedded_dg':
                         self.wrapper.subwrappers.update({field:EmbeddedDGWrapper(self, suboption)})
                     elif suboption.name == "recovered":
@@ -136,9 +116,9 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                     else:
                         raise RuntimeError(
                         f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
-                print(self.wrapper)
-                print(self.wrapper.suboptions)
-                print(self.wrapper.subwrappers)
+                #print(self.wrapper)
+                #print(self.wrapper.suboptions)
+                #print(self.wrapper.subwrappers)
                 #self.wrapper_name = 'mixed'
             else:
                 self.wrapper_name = options.name
@@ -219,35 +199,54 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 # Set these up with ?
                 
                 # Give more than one fs?
-                fields = []
+                #fields = []
                 
                 print(self.wrapper.wrapper_spaces)
                 
                 for field, subwrapper in self.wrapper.subwrappers.items():
-                    # Set up field idxs here.
-                    print(field)
-                    print(subwrapper)
-                    self.wrapper.subwrappers[field].idx = equation.field_names.index(field)
+                    field_idx = equation.field_names.index(field)
+                    self.wrapper.subwrappers[field].idx = field_idx
                     self.wrapper.subwrappers[field].mixed_options = True
                     
                     # Store the original space of the tracer
                     self.wrapper.subwrappers[field].tracer_fs = self.equation.spaces[equation.field_names.index(field)]
+                    
                     self.wrapper.subwrappers[field].setup()
                     
                     # Update the function space to that needed by the wrapper
-                    self.wrapper.wrapper_spaces.update({field: self.wrapper.subwrappers[field].function_space})
-                    fields.append(self.wrapper.subwrappers[field].function_space)
-                    # Append the new function space from the wrapper
-                    #print('self.tracer_fs is', self.wrapper.subwrappers[field].tracer_fs)
-                print(self.wrapper.wrapper_spaces)
-                #Setup the mixed wrapper
-                #yo_spaces = []
-                #for _, space_name in self.wrapper.wrapper_spaces.items():
-                #    yo_spaces.append(self.domain.spaces(space_name))
-                #print(yo_spaces)
-                self.wrapper.wrapper_spaces = fields
+                    self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space
+                    
+                    # Replace with new test function
+                    # STILL TO SORT OUT.
+                    #new_test = TestFunction(self.wrapper.subwrappers[field].test_space)
+                    
+                    #self.residual = self.residual.label_map(
+                    #    all_terms,
+                    #    map_if_true=replace_test_function(new_test, old_idx=field_idx))
+                    
+                
                 self.wrapper.setup()
-                pass
+                
+                # Check if mixed function spcae has changed:
+                if self.wrapper.function_space == equation.function_space:
+                    print('same')
+                else:
+                    print('different')
+                    
+                self.fs = self.wrapper.function_space
+                
+                # Or replace test functions here??
+                new_test = TestFunction(self.wrapper.test_space)
+                
+                self.residual = self.residual.label_map(
+                        all_terms,
+                        map_if_true=replace_test_function(new_test))
+                
+                # Only call this if with SUPG, so should put this into the 
+                # previous section
+                self.residual = self.wrapper.label_terms(self.residual)
+                
+                
             else:
                 self.wrapper.setup()
                 self.fs = self.wrapper.function_space
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index d4fa6d4dd..0e6c655d5 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -33,7 +33,6 @@ def __init__(self, time_discretisation, wrapper_options):
         self.time_discretisation = time_discretisation
         self.options = wrapper_options
         self.solver_parameters = None
-        #self.field_name = None
         self.idx = None
         self.mixed_options = False
         self.tracer_fs = None
@@ -93,6 +92,7 @@ def setup(self):
             #print('self.tracer_fs is', self.tracer_fs)
             #print('self.time_discretisation.fs is', self.time_discretisation.fs)
             original_space = self.tracer_fs
+            #original_space = self.time_discretisation.fs
         else:
             original_space = self.time_discretisation.fs
 
@@ -114,6 +114,7 @@ def setup(self):
 
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
+        
         if self.mixed_options == True:
             self.x_projected = Function(self.tracer_fs)
         elif self.time_discretisation.idx is None:
@@ -182,7 +183,8 @@ def setup(self):
         equation = self.time_discretisation.equation
         
         if self.mixed_options == True:
-            original_space = self.tracer_fs
+            #original_space = self.tracer_fs
+            original_space = self.time_discretisation.fs
         else:
             original_space = self.time_discretisation.fs
 
@@ -416,9 +418,12 @@ def __init__(self, equation, suboptions):
         #print(equation.active_tracers)
         #print(equation.active_tracers[0])
         #print(equation.active_tracers[1])
-        print(equation.space_names)
+        #print(equation.space_names)
+        #print(equation.spaces)
         
-        self.wrapper_spaces = equation.space_names
+        #self.wrapper_spaces = equation.space_names
+        self.wrapper_spaces = equation.spaces
+        print(len(equation.spaces))
         
         #self.x_in = Function(equation.function_space)
         #self.x_out = Function(equation.function_space)
@@ -432,8 +437,8 @@ def __init__(self, equation, suboptions):
         self.subwrappers = {}
         
         for field, suboption in suboptions.items():
-                print(field)
-                print(suboption)
+                #print(field)
+                #print(suboption)
                 
                 # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
@@ -465,32 +470,45 @@ def __init__(self, equation, suboptions):
     def setup(self):
         # This is done in the suboption wrappers themselves
         # Or, determine the new mixed function space
+        
+        #Loop over all active variables?
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space) 
-        
-        
+        self.test_space = self.function_space
         #pass
     
     def pre_apply(self, x_in):
         """
         Perform the pre-applications from all subwrappers
         """
+        #self.x_in.assign(x_in)
+        self.x_in = x_in
 
         for _, subwrapper in self.subwrappers.items():
             print(subwrapper)
             print('pre')
             field = x_in.subfunctions[subwrapper.idx]
             subwrapper.pre_apply(field)
-            print('success')
+            #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
+            
+            x_in_sub = self.x_in.subfunctions[subwrapper.idx]
+            #x_in_sub.assign(subwrapper.x_in)
+            x_in_sub = subwrapper.x_in
+            #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
             
-    def post_apply(self, x_in):
+            
+    def post_apply(self, x_out):
         """
         Perform the post-applications from all subwrappers
         """
+        x_out.assign(self.x_out)
 
         for _, subwrapper in self.subwrappers.items():
             print(subwrapper)
             print('post')
-            field = x_in.subfunctions[subwrapper.idx]
-            subwrapper.post_apply(field)
\ No newline at end of file
+            field = x_out.subfunctions[subwrapper.idx]
+            subwrapper.post_apply(field)
+            
+            x_out_sub = x_out.subfunctions[subwrapper.idx]
+            x_out_sub.assign(subwrapper.x_out)
\ No newline at end of file

From 5f05367a13b13fd92fda235e3c92295140447881 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Thu, 11 Jan 2024 04:09:03 +0000
Subject: [PATCH 08/38] more changes to test functions for multiple wrappers

---
 gusto/time_discretisation.py | 16 ++++++++--------
 gusto/wrappers.py            |  4 ++--
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 90f73a502..e77dff92e 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -216,14 +216,14 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space
                     
-                    # Replace with new test function
-                    # STILL TO SORT OUT.
-                    #new_test = TestFunction(self.wrapper.subwrappers[field].test_space)
-                    
-                    #self.residual = self.residual.label_map(
-                    #    all_terms,
-                    #    map_if_true=replace_test_function(new_test, old_idx=field_idx))
                     
+                    # Replace the test function space
+                    # This currently won't work: supg test 
+                    # is a function, not a space
+                    if self.wrapper.suboptions[field].name == "supg":
+                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test
+                    else:
+                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test_space
                 
                 self.wrapper.setup()
                 
@@ -244,7 +244,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 
                 # Only call this if with SUPG, so should put this into the 
                 # previous section
-                self.residual = self.wrapper.label_terms(self.residual)
+                # self.residual = self.wrapper.label_terms(self.residual)
                 
                 
             else:
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 0e6c655d5..052ddbc2e 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -423,6 +423,7 @@ def __init__(self, equation, suboptions):
         
         #self.wrapper_spaces = equation.space_names
         self.wrapper_spaces = equation.spaces
+        self.test_spaces = equation.spaces
         print(len(equation.spaces))
         
         #self.x_in = Function(equation.function_space)
@@ -475,8 +476,7 @@ def setup(self):
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space) 
-        self.test_space = self.function_space
-        #pass
+        self.test_space = MixedFunctionSpace(self.test_spaces)
     
     def pre_apply(self, x_in):
         """

From 5c13db13af2e641d5cb74390f0a020ef170cf8b4 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 16 Jan 2024 22:26:08 +0000
Subject: [PATCH 09/38] more changes to mixed options with replacing test
 functions

---
 gusto/time_discretisation.py                  |  93 +++++++---------
 gusto/wrappers.py                             | 105 ++++++++++--------
 .../transport/test_mixed_fs_options.py        |  10 +-
 3 files changed, 102 insertions(+), 106 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index e77dff92e..8410a0677 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -21,7 +21,7 @@
 
 from gusto.configuration import EmbeddedDGOptions, RecoveryOptions
 from gusto.labels import (time_derivative, prognostic, physics_label,
-                          implicit, explicit)
+                          implicit, explicit, transport)
 from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.wrappers import *
 
@@ -43,15 +43,6 @@ def new_apply(self, x_out, x_in):
 
                 self.wrapper.pre_apply(x_in)
                 original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
-                
-                #if type(self.wrapper) == MixedOptions:
-                    #print('x_out', x_out)
-                    #print('self.wrapper.x_out', self.wrapper.x_out)
-                    #original_apply(self, x_out, x_in)
-                #    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
-                #else:
-                #    original_apply(self, self.wrapper.x_out, self.wrapper.x_in)
-                
                 self.wrapper.post_apply(x_out)
 
             return new_apply(self, x_out, x_in)
@@ -97,12 +88,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
 
         if options is not None:
             if type(options) == MixedOptions:
-                #print('jahjah')
-                #print(options)
                 self.wrapper = options
-                #self.subwrappers = {}
-                # Or do I need to initialise everything here
-                # like is done for a single wrapper?
+                
                 for field, suboption in self.wrapper.suboptions.items():
                     print(field)
                     print(suboption)
@@ -116,10 +103,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                     else:
                         raise RuntimeError(
                         f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
-                #print(self.wrapper)
-                #print(self.wrapper.suboptions)
-                #print(self.wrapper.subwrappers)
-                #self.wrapper_name = 'mixed'
+                        
             else:
                 self.wrapper_name = options.name
                 if self.wrapper_name == "embedded_dg":
@@ -194,58 +178,61 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         if self.wrapper is not None:
             if type(self.wrapper) == MixedOptions:
-            #if self.wrapper_name == 'mixed':
-                # Subwrappers are defined.
-                # Set these up with ?
-                
-                # Give more than one fs?
-                #fields = []
-                
-                print(self.wrapper.wrapper_spaces)
+            
+                #print(self.wrapper.wrapper_spaces)
                 
                 for field, subwrapper in self.wrapper.subwrappers.items():
                     field_idx = equation.field_names.index(field)
-                    self.wrapper.subwrappers[field].idx = field_idx
-                    self.wrapper.subwrappers[field].mixed_options = True
+                    
+                    subwrapper.idx = field_idx
+                    subwrapper.mixed_options = True
                     
                     # Store the original space of the tracer
-                    self.wrapper.subwrappers[field].tracer_fs = self.equation.spaces[equation.field_names.index(field)]
+                    subwrapper.tracer_fs = self.equation.spaces[field_idx]
                     
-                    self.wrapper.subwrappers[field].setup()
+                    subwrapper.setup()
                     
                     # Update the function space to that needed by the wrapper
-                    self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space
+                    self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
                     
+                    # Store test space?
+                    #self.wrapper.test_spaces[field_idx] = subwrapper.function_space
                     
                     # Replace the test function space
-                    # This currently won't work: supg test 
-                    # is a function, not a space
                     if self.wrapper.suboptions[field].name == "supg":
-                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test
+                        new_test = subwrapper.test
                     else:
-                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test_space
+                        new_test = TestFunction(subwrapper.test_space)
+                        
+                    self.residual = self.residual.label_map(
+                        lambda t: t.has_label(transport) and t.get(prognostic) == field,
+                        map_if_true=replace_test_function(new_test, old_idx=field_idx))
+                        
+                    self.residual = subwrapper.label_terms(self.residual)
+
+                    # Currently can only use one set of solver parameters ...
+                    if self.solver_parameters is None:
+                        self.solver_parameters = subwrapper.solver_parameters
                 
                 self.wrapper.setup()
-                
-                # Check if mixed function spcae has changed:
-                if self.wrapper.function_space == equation.function_space:
-                    print('same')
-                else:
-                    print('different')
                     
                 self.fs = self.wrapper.function_space
                 
-                # Or replace test functions here??
-                new_test = TestFunction(self.wrapper.test_space)
-                
-                self.residual = self.residual.label_map(
-                        all_terms,
-                        map_if_true=replace_test_function(new_test))
-                
-                # Only call this if with SUPG, so should put this into the 
-                # previous section
-                # self.residual = self.wrapper.label_terms(self.residual)
-                
+                #new_test_mixed = TestFunction(self.fs)
+
+                #for field, subwrapper in self.wrapper.subwrappers.items():
+                #    field_idx = equation.field_names.index(field)
+
+                #   if self.wrapper.suboptions[field].name == "supg":
+                #        new_test_mixed[field_idx] = subwrapper.test
+                        
+                #    self.residual = self.residual.label_map(
+                #        lambda t: t.has_label(transport) and t.get(prognostic) == field,
+                #        map_if_true=replace_test_function(new_test_mixed[field_idx], old_idx=field_idx))
+                        
+                #    self.residual = subwrapper.label_terms(self.residual)
+                    
+                    
                 
             else:
                 self.wrapper.setup()
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 052ddbc2e..4508ffde3 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -414,39 +414,21 @@ def __init__(self, equation, suboptions):
         Raises:
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
-        #print(equation.function_space)
-        #print(equation.active_tracers)
-        #print(equation.active_tracers[0])
-        #print(equation.active_tracers[1])
-        #print(equation.space_names)
-        #print(equation.spaces)
-        
-        #self.wrapper_spaces = equation.space_names
         self.wrapper_spaces = equation.spaces
         self.test_spaces = equation.spaces
-        print(len(equation.spaces))
-        
-        #self.x_in = Function(equation.function_space)
-        #self.x_out = Function(equation.function_space)
-        
-        print('Initialising MixedOptions')
-        
+        self.field_names = equation.field_names
         self.suboptions = suboptions
-        
-        #print(suboptions.items())
-        
         self.subwrappers = {}
         
         for field, suboption in suboptions.items():
-                #print(field)
-                #print(suboption)
                 
                 # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
-                else:
+                #else:
                     # check that a valid wrapper has been given
-                    wrapper_name = suboption.name
+                
+               #    wrapper_name = suboption.name
                     #print(wrapper_name)
                     
                     # Extract the space?
@@ -469,10 +451,8 @@ def __init__(self, equation, suboptions):
  #                   self.subwrapper.x_in = Function(self.suboptions[field].fs)
     
     def setup(self):
-        # This is done in the suboption wrappers themselves
-        # Or, determine the new mixed function space
         
-        #Loop over all active variables?
+        # Compute the new mixed function space
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space) 
@@ -480,35 +460,64 @@ def setup(self):
     
     def pre_apply(self, x_in):
         """
-        Perform the pre-applications from all subwrappers
+        Perform the pre-applications for all fields
+        with an associated subwrapper.
         """
-        #self.x_in.assign(x_in)
-        self.x_in = x_in
-
-        for _, subwrapper in self.subwrappers.items():
-            print(subwrapper)
-            print('pre')
-            field = x_in.subfunctions[subwrapper.idx]
-            subwrapper.pre_apply(field)
-            #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
             
-            x_in_sub = self.x_in.subfunctions[subwrapper.idx]
-            #x_in_sub.assign(subwrapper.x_in)
-            x_in_sub = subwrapper.x_in
-            #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
+        for field_name in self.field_names:
+            #print(field_name)
+            #print(self.subwrappers)
+            
+            field_idx = self.field_names.index(field_name)
+            #print(field_idx)
+            
+            field = x_in.subfunctions[field_idx]
+            x_in_sub = self.x_in.subfunctions[field_idx]
+            
+            if field_name in self.subwrappers:
+                subwrapper = self.subwrappers[field_name]
+                print(subwrapper)
+                print('pre')
+                #field = x_in.subfunctions[subwrapper.idx]
+                subwrapper.pre_apply(field)
+                #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
+            
+                #x_in_sub = self.x_in.subfunctions[subwrapper.idx]
+                x_in_sub.assign(subwrapper.x_in)
+                #x_in_sub = subwrapper.x_in
+                #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
+            else:
+                x_in_sub.assign(field)        
             
             
     def post_apply(self, x_out):
         """
-        Perform the post-applications from all subwrappers
+        Perform the post-applications for all fields
+        with an associated subwrapper.
         """
-        x_out.assign(self.x_out)
+        #x_out.assign(self.x_out)
 
-        for _, subwrapper in self.subwrappers.items():
-            print(subwrapper)
-            print('post')
-            field = x_out.subfunctions[subwrapper.idx]
-            subwrapper.post_apply(field)
+        #for _, subwrapper in self.subwrappers.items():
+        #    print(subwrapper)
+        #    print('post')
+        #    field = x_out.subfunctions[subwrapper.idx]
+        #    subwrapper.post_apply(field)
             
-            x_out_sub = x_out.subfunctions[subwrapper.idx]
-            x_out_sub.assign(subwrapper.x_out)
\ No newline at end of file
+        #    x_out_sub = x_out.subfunctions[subwrapper.idx]
+        #    x_out_sub.assign(subwrapper.x_out)
+            
+        for field_name in self.field_names:
+        
+            field_idx = self.field_names.index(field_name)
+            
+            field = self.x_out.subfunctions[field_idx]
+            x_out_sub = x_out.subfunctions[field_idx]
+            
+            if field_name in self.subwrappers:
+                subwrapper = self.subwrappers[field_name]
+                print(subwrapper)
+                print('post')
+                subwrapper.post_apply(field)
+                x_out_sub.assign(subwrapper.x_out)
+            else:
+                x_out_sub.assign(field)  
\ No newline at end of file
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 699618a64..5e6c4739b 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -64,25 +64,23 @@ def setup_limiters(dirname, space_A, space_B):
     
     # Tracer B spaces
     if space_B == 'DG':
+        space_B_string = 'DG'
         if degree == 0:
             VB = domain.spaces('DG')
             VCG1_B = FunctionSpace(mesh, 'CG', 1)
             VDG1_B = domain.spaces('DG1_equispaced')
-            space_B_string = 'DG'
         elif degree == 1:
             VB = domain.spaces('DG')
-            space_B_string = 'DG'
         else:
             raise NotImplementedError
     elif space_B == 'Vtheta':
+        space_B_string = 'theta'
         if degree == 0: 
             VB = domain.spaces('theta')
             VCG1_B = FunctionSpace(mesh, 'CG', 1)
             VDG1_B = domain.spaces('DG1_equispaced')
-            space_B_string = 'theta'
         elif degree == 1:
             VB = domain.spaces('theta')
-            space_B_string = 'theta'
         else:
             raise NotImplementedError
     else:
@@ -178,6 +176,9 @@ def setup_limiters(dirname, space_A, space_B):
     # DG Upwind transport for both tracers:
     transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
     
+    # Need to give SUPG options to the above, if using supg ...
+    # Need to test SUPG here!
+    
     # Build time stepper
     stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method)
 
@@ -279,7 +280,6 @@ def setup_limiters(dirname, space_A, space_B):
 
 @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
                                    'DG1', 'DG1_equispaced'])
-#@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced'])
 # It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 

From 806a49ef5701c33660d6916ac48622b6f7194291 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 17 Jan 2024 21:38:21 +0000
Subject: [PATCH 10/38] mixed options works for embedded dg and recovery

---
 gusto/equations.py                            |  2 -
 gusto/limiters.py                             |  2 -
 gusto/time_discretisation.py                  | 50 +++++++-----
 gusto/wrappers.py                             | 81 +++++--------------
 .../transport/test_mixed_fs_options.py        | 21 ++++-
 5 files changed, 70 insertions(+), 86 deletions(-)

diff --git a/gusto/equations.py b/gusto/equations.py
index c30124c73..1f3b90afd 100644
--- a/gusto/equations.py
+++ b/gusto/equations.py
@@ -426,8 +426,6 @@ def add_tracers_to_prognostics(self, domain, active_tracers):
                 else:
                     self.space_names[tracer.name] = tracer.space
                 self.spaces.append(domain.spaces(tracer.space))
-                #print('self.spaces is', self.spaces)
-                print(len(self.spaces))
             else:
                 raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}')
 
diff --git a/gusto/limiters.py b/gusto/limiters.py
index 16a2d3b40..a39c6688d 100644
--- a/gusto/limiters.py
+++ b/gusto/limiters.py
@@ -203,8 +203,6 @@ def __init__(self, equation, sublimiters):
                 raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
             else:
                 self.sublimiters[field].idx = equation.field_names.index(field)
-                
-                # Give the sublimiter a subspace?
 
     def apply(self, fields):
         """
diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 8410a0677..d883ae216 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -10,7 +10,7 @@
 import numpy as np
 
 from firedrake import (
-    Function, TestFunction, NonlinearVariationalProblem,
+    Function, TestFunction, TestFunctions, NonlinearVariationalProblem,
     NonlinearVariationalSolver, DirichletBC, split, Constant
 )
 from firedrake.fml import (
@@ -199,39 +199,53 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     #self.wrapper.test_spaces[field_idx] = subwrapper.function_space
                     
                     # Replace the test function space
-                    if self.wrapper.suboptions[field].name == "supg":
-                        new_test = subwrapper.test
-                    else:
-                        new_test = TestFunction(subwrapper.test_space)
+                    #if self.wrapper.suboptions[field].name == "supg":
+                    #    new_test = subwrapper.test
+                    #else:
+                    #    new_test = TestFunction(subwrapper.test_space)
                         
-                    self.residual = self.residual.label_map(
-                        lambda t: t.has_label(transport) and t.get(prognostic) == field,
-                        map_if_true=replace_test_function(new_test, old_idx=field_idx))
+                    #self.residual = self.residual.label_map(
+                    #    lambda t: t.has_label(transport) and t.get(prognostic) == field,
+                    #    map_if_true=replace_test_function(new_test, old_idx=field_idx))
                         
-                    self.residual = subwrapper.label_terms(self.residual)
-
-                    # Currently can only use one set of solver parameters ...
-                    if self.solver_parameters is None:
-                        self.solver_parameters = subwrapper.solver_parameters
+                    #self.residual = subwrapper.label_terms(self.residual)
                 
                 self.wrapper.setup()
                     
                 self.fs = self.wrapper.function_space
                 
-                #new_test_mixed = TestFunction(self.fs)
+                new_test_mixed = TestFunctions(self.fs)
 
-                #for field, subwrapper in self.wrapper.subwrappers.items():
-                #    field_idx = equation.field_names.index(field)
+                # If use supg, then change the required test function
 
-                #   if self.wrapper.suboptions[field].name == "supg":
-                #        new_test_mixed[field_idx] = subwrapper.test
+                # Replace one-by-one in this case?
+
+                #for field, subwrapper in self.wrapper.subwrappers.items():
+                #    if self.wrapper.suboptions[field].name == "supg":
+                #        field_idx = equation.field_names.index(field)
+                #        test_list = list(new_test_mixed)
+                #        test_list[field_idx] = subwrapper.test
+                #        new_test_mixed = tuple(test_list)
                         
+                # Change test functions with the new space:
+                #for field, subwrapper in self.wrapper.subwrappers.items():        
                 #    self.residual = self.residual.label_map(
                 #        lambda t: t.has_label(transport) and t.get(prognostic) == field,
                 #        map_if_true=replace_test_function(new_test_mixed[field_idx], old_idx=field_idx))
                         
                 #    self.residual = subwrapper.label_terms(self.residual)
                     
+                # Or, all-at-once:
+                #self.residual = self.residual.label_map(
+                #        lambda t: t.has_label(transport),
+                #        map_if_true=replace_test_function(new_test_mixed))
+                        
+                self.residual = self.residual.label_map(
+                        all_terms,
+                        map_if_true=replace_test_function(new_test_mixed))
+                        
+                
+                    
                     
                 
             else:
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 4508ffde3..956f409d2 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -89,10 +89,7 @@ def setup(self):
         equation = self.time_discretisation.equation
         
         if self.mixed_options == True:
-            #print('self.tracer_fs is', self.tracer_fs)
-            #print('self.time_discretisation.fs is', self.time_discretisation.fs)
             original_space = self.tracer_fs
-            #original_space = self.time_discretisation.fs
         else:
             original_space = self.time_discretisation.fs
 
@@ -207,19 +204,12 @@ def setup(self):
 
         if self.mixed_options == True:
             self.x_in_tmp = Function(self.tracer_fs)
-            #self.x_in_tmp = Function(equation.spaces[self.idx])
         else:
             self.x_in_tmp = Function(self.time_discretisation.fs)
             
-        #print(self.function_space)
-        #print(self.time_discretisation.fs)
-            
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
         
-        #self.x_out = Function(equation.function_space)
-        #self.x_out = Function(self.time_discretisation.fs)
-        
         if self.mixed_options == True:
             self.x_projected = Function(self.tracer_fs)
         elif self.time_discretisation.idx is None:
@@ -306,9 +296,23 @@ def setup(self):
             'SUPG wrapper can only be used with SUPG Options'
 
         domain = self.time_discretisation.domain
+        
+        #if self.mixed_options == True:
+        #    self.function_space = self.tracer_fs
+        #else:
+        #    self.function_space = self.time_discretisation.fs
+            
         self.function_space = self.time_discretisation.fs
-        self.test_space = self.function_space
-        self.x_out = Function(self.function_space)
+        
+        if self.mixed_options == True:
+            self.test_space = self.tracer_fs
+            self.x_out = Function(self.tracer_fs)
+        else:
+            self.test_space = self.function_space
+            self.x_out = Function(self.function_space)
+        
+        #self.test_space = self.function_space
+        #self.x_out = Function(self.function_space)
 
         # -------------------------------------------------------------------- #
         # Work out SUPG parameter
@@ -415,40 +419,15 @@ def __init__(self, equation, suboptions):
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
         self.wrapper_spaces = equation.spaces
-        self.test_spaces = equation.spaces
+        #self.test_spaces = equation.spaces
         self.field_names = equation.field_names
         self.suboptions = suboptions
         self.subwrappers = {}
         
         for field, suboption in suboptions.items():
-                
                 # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
-                #else:
-                    # check that a valid wrapper has been given
-                
-               #    wrapper_name = suboption.name
-                    #print(wrapper_name)
-                    
-                    # Extract the space?
-                    
-                    
-                    #if wrapper_name == "embedded_dg":
-                    #    self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)})
-                    #elif wrapper_name == "recovered":
-                    #    self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption)
-                    #elif wrapper_name == "supg":
-                    #    self.suboptions[field].subwrapper = SUPGWrapper(self, suboption)
-                    #else:
-                    #    raise RuntimeError(
-                    #    f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
-                    
-                    #Initialise the wrapper and associate with a field:
-                    #self.suboptions[field].wrapper_name = suboption.name
-                    #self.suboptions[field].idx = equation.field_names.index(field)
-                    #self.suboptions[field].fs = equation.field_names.function_space(field)
- #                   self.subwrapper.x_in = Function(self.suboptions[field].fs)
     
     def setup(self):
         
@@ -456,7 +435,7 @@ def setup(self):
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space) 
-        self.test_space = MixedFunctionSpace(self.test_spaces)
+        #self.test_space = MixedFunctionSpace(self.test_spaces)
     
     def pre_apply(self, x_in):
         """
@@ -476,16 +455,8 @@ def pre_apply(self, x_in):
             
             if field_name in self.subwrappers:
                 subwrapper = self.subwrappers[field_name]
-                print(subwrapper)
-                print('pre')
-                #field = x_in.subfunctions[subwrapper.idx]
                 subwrapper.pre_apply(field)
-                #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
-            
-                #x_in_sub = self.x_in.subfunctions[subwrapper.idx]
                 x_in_sub.assign(subwrapper.x_in)
-                #x_in_sub = subwrapper.x_in
-                #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in
             else:
                 x_in_sub.assign(field)        
             
@@ -495,16 +466,6 @@ def post_apply(self, x_out):
         Perform the post-applications for all fields
         with an associated subwrapper.
         """
-        #x_out.assign(self.x_out)
-
-        #for _, subwrapper in self.subwrappers.items():
-        #    print(subwrapper)
-        #    print('post')
-        #    field = x_out.subfunctions[subwrapper.idx]
-        #    subwrapper.post_apply(field)
-            
-        #    x_out_sub = x_out.subfunctions[subwrapper.idx]
-        #    x_out_sub.assign(subwrapper.x_out)
             
         for field_name in self.field_names:
         
@@ -515,9 +476,7 @@ def post_apply(self, x_out):
             
             if field_name in self.subwrappers:
                 subwrapper = self.subwrappers[field_name]
-                print(subwrapper)
-                print('post')
-                subwrapper.post_apply(field)
-                x_out_sub.assign(subwrapper.x_out)
+                subwrapper.x_out.assign(field)
+                subwrapper.post_apply(x_out_sub)
             else:
                 x_out_sub.assign(field)  
\ No newline at end of file
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 5e6c4739b..301c260e5 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -59,6 +59,12 @@ def setup_limiters(dirname, space_A, space_B):
     elif space_A == 'Vtheta_degree_1':
         VA = domain.spaces('theta')
         space_A_string = 'theta'
+    elif space_A == 'HDiv':
+        VA = domain.spaces('HDiv')
+        space_A_string = 'HDiv'
+    #elif space_A == 'CG1'
+        #VA = FunctionSpace(mesh, 'CG', 1)
+        #space_A_string = 'CG'
     else:
         raise NotImplementedError
     
@@ -131,7 +137,9 @@ def setup_limiters(dirname, space_A, space_B):
     elif space_A == 'Vtheta_degree_1':
         suboptions.update({'tracerA': EmbeddedDGOptions()})
         sublimiters.update({'tracerA': ThetaLimiter(VA)})
-
+    elif space_A == 'HDiv':
+        ibp_A = IntegrateByParts.TWICE
+        suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)})
     else:
         raise NotImplementedError
         
@@ -174,7 +182,13 @@ def setup_limiters(dirname, space_A, space_B):
     #transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
     
     # DG Upwind transport for both tracers:
-    transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
+    if space_A == 'HDiv':
+        transport_method_A = DGUpwind(eqn, 'tracerA', ibp=ibp_A)
+    else: 
+        transport_method_A = DGUpwind(eqn, 'tracerA')
+    transport_method_B = DGUpwind(eqn, 'tracerB')
+    
+    transport_method = [transport_method_A, transport_method_B]
     
     # Need to give SUPG options to the above, if using supg ...
     # Need to test SUPG here!
@@ -278,8 +292,9 @@ def setup_limiters(dirname, space_A, space_B):
     return stepper, tmax, true_fieldA, true_fieldB
 
 
+#@pytest.mark.parametrize('space_A', ['HDiv'])#, 
 @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-                                   'DG1', 'DG1_equispaced'])
+                                     'DG1', 'DG1_equispaced', 'HDiv'])
 # It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 

From 2be135486d0ef1d29e0648d25658bcff3967027b Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 02:09:55 +0000
Subject: [PATCH 11/38] moved DG1-DG1 equispaced test from test_limiters to
 test_mixed_fs_options

---
 gusto/equations.py                            |   1 -
 gusto/time_discretisation.py                  |  50 +++----
 gusto/wrappers.py                             |  11 +-
 integration-tests/transport/test_limiters.py  | 124 ++++--------------
 .../transport/test_mixed_fs_options.py        |  18 ++-
 5 files changed, 67 insertions(+), 137 deletions(-)

diff --git a/gusto/equations.py b/gusto/equations.py
index 1f3b90afd..bd61e795b 100644
--- a/gusto/equations.py
+++ b/gusto/equations.py
@@ -542,7 +542,6 @@ def __init__(self, domain, active_tracers, Vu=None):
         # Build finite element spaces
         self.spaces = []
 
-
         # Add active tracers to the list of prognostics
         if active_tracers is None:
             active_tracers = []
diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index d883ae216..0ce9dc68e 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -21,7 +21,7 @@
 
 from gusto.configuration import EmbeddedDGOptions, RecoveryOptions
 from gusto.labels import (time_derivative, prognostic, physics_label,
-                          implicit, explicit, transport)
+                          implicit, explicit)
 from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.wrappers import *
 
@@ -89,21 +89,21 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         if options is not None:
             if type(options) == MixedOptions:
                 self.wrapper = options
-                
+
                 for field, suboption in self.wrapper.suboptions.items():
                     print(field)
                     print(suboption)
-                    
+
                     if suboption.name == 'embedded_dg':
-                        self.wrapper.subwrappers.update({field:EmbeddedDGWrapper(self, suboption)})
+                        self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)})
                     elif suboption.name == "recovered":
-                        self.wrapper.subwrappers.update({field:RecoveryWrapper(self, suboption)})
+                        self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)})
                     elif suboption.name == "supg":
-                        self.wrapper.subwrappers.update({field:SUPGWrapper(self, suboption)})
+                        self.wrapper.subwrappers.update({field: SUPGWrapper(self, suboption)})
                     else:
                         raise RuntimeError(
-                        f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
-                        
+                            f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
+
             else:
                 self.wrapper_name = options.name
                 if self.wrapper_name == "embedded_dg":
@@ -178,23 +178,21 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         if self.wrapper is not None:
             if type(self.wrapper) == MixedOptions:
-            
-                #print(self.wrapper.wrapper_spaces)
-                
+
                 for field, subwrapper in self.wrapper.subwrappers.items():
                     field_idx = equation.field_names.index(field)
-                    
+
                     subwrapper.idx = field_idx
                     subwrapper.mixed_options = True
-                    
+
                     # Store the original space of the tracer
                     subwrapper.tracer_fs = self.equation.spaces[field_idx]
-                    
+
                     subwrapper.setup()
-                    
+
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
-                    
+
                     # Store test space?
                     #self.wrapper.test_spaces[field_idx] = subwrapper.function_space
                     
@@ -219,13 +217,21 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 # If use supg, then change the required test function
 
                 # Replace one-by-one in this case?
+                print(type(new_test_mixed))
 
-                #for field, subwrapper in self.wrapper.subwrappers.items():
-                #    if self.wrapper.suboptions[field].name == "supg":
-                #        field_idx = equation.field_names.index(field)
-                #        test_list = list(new_test_mixed)
-                #        test_list[field_idx] = subwrapper.test
-                #        new_test_mixed = tuple(test_list)
+                for field, subwrapper in self.wrapper.subwrappers.items():
+                    if self.wrapper.suboptions[field].name == "supg":
+                        field_idx = equation.field_names.index(field)
+                        test_list = list(new_test_mixed)
+                        test_list[field_idx] = subwrapper.test
+                        new_test_mixed2 = tuple(test_list)
+                        
+                print(type(new_test_mixed2))
+                        
+                if new_test_mixed == new_test_mixed2:
+                    print('same')
+                else:
+                    print('different')
                         
                 # Change test functions with the new space:
                 #for field, subwrapper in self.wrapper.subwrappers.items():        
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 956f409d2..900224733 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -180,8 +180,7 @@ def setup(self):
         equation = self.time_discretisation.equation
         
         if self.mixed_options == True:
-            #original_space = self.tracer_fs
-            original_space = self.time_discretisation.fs
+            original_space = self.tracer_fs
         else:
             original_space = self.time_discretisation.fs
 
@@ -302,7 +301,9 @@ def setup(self):
         #else:
         #    self.function_space = self.time_discretisation.fs
             
+        # But this will create some problems when defining the new mixed space...
         self.function_space = self.time_discretisation.fs
+        #self.function_space = self.tracer_fs
         
         if self.mixed_options == True:
             self.test_space = self.tracer_fs
@@ -444,12 +445,8 @@ def pre_apply(self, x_in):
         """
             
         for field_name in self.field_names:
-            #print(field_name)
-            #print(self.subwrappers)
-            
+        
             field_idx = self.field_names.index(field_name)
-            #print(field_idx)
-            
             field = x_in.subfunctions[field_idx]
             x_in_sub = self.x_in.subfunctions[field_idx]
             
diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py
index 9853b11c3..94d83ae92 100644
--- a/integration-tests/transport/test_limiters.py
+++ b/integration-tests/transport/test_limiters.py
@@ -50,32 +50,17 @@ def setup_limiters(dirname, space):
         VDG1 = domain.spaces('DG1_equispaced')
     elif space == 'Vtheta_degree_1':
         V = domain.spaces('theta')
-    elif space == 'mixed_FS':
-        V = domain.spaces('HDiv')
-        VA = domain.spaces('DG')
-        VB = domain.spaces('DG1_equispaced')
     else:
         raise NotImplementedError
 
     Vpsi = domain.spaces('H1')
 
     # Equation
-    if space == 'mixed_FS':
-        tracerA = ActiveTracer(name='tracerA', space='DG',
-                               variable_type=TracerVariableType.mixing_ratio,
-                               transport_eqn=TransportEquationType.advective)
-        tracerB = ActiveTracer(name='tracerB', space='DG1_equispaced',
-                               variable_type=TracerVariableType.mixing_ratio,
-                               transport_eqn=TransportEquationType.advective)
-        tracers = [tracerA, tracerB]
-        eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V)
-        output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1,
-                                  dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB'])
-    else:
-        eqn = AdvectionEquation(domain, V, 'tracer')
-        output = OutputParameters(dirname=dirname+'/limiters',
-                                  dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer'])
+    eqn = AdvectionEquation(domain, V, 'tracer')
 
+    # I/O
+    output = OutputParameters(dirname=dirname+'/limiters',
+                              dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer'])
     io = IO(domain, output)
 
     # ------------------------------------------------------------------------ #
@@ -99,19 +84,10 @@ def setup_limiters(dirname, space):
     elif space == 'Vtheta_degree_1':
         opts = EmbeddedDGOptions()
         transport_schemes = SSPRK3(domain, options=opts, limiter=ThetaLimiter(V))
-
-    elif space == 'mixed_FS':
-        sublimiters = {'tracerA': DG1Limiter(domain.spaces('DG')),
-                       'tracerB': VertexBasedLimiter(domain.spaces('DG1_equispaced'))}
-        MixedLimiter = MixedFSLimiter(eqn, sublimiters)
-        transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
     else:
         raise NotImplementedError
 
-    if space == 'mixed_FS':
-        transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
-    else:
-        transport_method = DGUpwind(eqn, "tracer")
+    transport_method = DGUpwind(eqn, "tracer")
 
     # Build time stepper
     stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method)
@@ -120,14 +96,8 @@ def setup_limiters(dirname, space):
     # Initial condition
     # ------------------------------------------------------------------------ #
 
-    if space == 'mixed_FS':
-        tracerA_0 = stepper.fields('tracerA', space=VA)
-        tracerB_0 = stepper.fields('tracerB', space=VB)
-        true_fieldA = stepper.fields('true_tracerA', space=VA)
-        true_fieldB = stepper.fields('true_tracerB', space=VB)
-    else:
-        tracer0 = stepper.fields('tracer', V)
-        true_field = stepper.fields('true_tracer', space=V)
+    tracer0 = stepper.fields('tracer', V)
+    true_field = stepper.fields('true_tracer', space=V)
 
     x, z = SpatialCoordinate(mesh)
 
@@ -176,17 +146,9 @@ def setup_limiters(dirname, space):
                              0.0)
 
         if i == 0:
-            if space == 'mixed_FS':
-                tracerA_0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
-                tracerB_0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
-            else:
-                tracer0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+            tracer0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
         elif i == 1:
-            if space == 'mixed_FS':
-                true_fieldA.interpolate(Constant(tracer_min) + expr_1 + expr_2)
-                true_fieldB.interpolate(Constant(tracer_min) + expr_1 + expr_2)
-            else:
-                true_field.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+            true_field.interpolate(Constant(tracer_min) + expr_1 + expr_2)
         else:
             raise ValueError
 
@@ -219,67 +181,29 @@ def setup_limiters(dirname, space):
     gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)])
     u.project(gradperp(psi))
 
-    if space == 'mixed_FS':
-        return stepper, tmax, true_fieldA, true_fieldB
-    else:
-        return stepper, tmax, true_field
+    return stepper, tmax, true_field
 
 
-@pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-                                   'DG1', 'DG1_equispaced', 'mixed_FS'])
+@pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1',
+                                   'DG0', 'DG1', 'DG1_equispaced'])
 def test_limiters(tmpdir, space):
 
     # Setup and run
     dirname = str(tmpdir)
-
-    if space == 'mixed_FS':
-        stepper, tmax, true_fieldA, true_fieldB = setup_limiters(dirname, space)
-    else:
-        stepper, tmax, true_field = setup_limiters(dirname, space)
-
+    stepper, tmax, true_field = setup_limiters(dirname, space)
     stepper.run(t=0, tmax=tmax)
+    final_field = stepper.fields('tracer')
 
-    tol = 1e-9
-
-    if space == 'mixed_FS':
-        final_fieldA = stepper.fields('tracerA')
-        final_fieldB = stepper.fields('tracerB')
+    # Check tracer is roughly in the correct place
+    assert norm(true_field - final_field) / norm(true_field) < 0.05, \
+        'Something appears to have gone wrong with transport of tracer using a limiter'
 
-        # Check tracer is roughly in the correct place
-        assert norm(true_fieldA - final_fieldA) / norm(true_fieldA) < 0.05, \
-            'Something is wrong with the DG space tracer using a mixed limiter'
-
-        # Check tracer is roughly in the correct place
-        assert norm(true_fieldB - final_fieldB) / norm(true_fieldB) < 0.05, \
-            'Something is wrong with the DG1 equispaced tracer using a mixed limiter'
-
-        # Check for no new overshoots in A
-        assert np.max(final_fieldA.dat.data) <= np.max(true_fieldA.dat.data) + tol, \
-            'Application of the DG space limiter in the mixed limiter has not prevented overshoots'
-
-        # Check for no new undershoots in A
-        assert np.min(final_fieldA.dat.data) >= np.min(true_fieldA.dat.data) - tol, \
-            'Application of the DG space limiter in the mixed limiter has not prevented undershoots'
-
-        # Check for no new overshoots in B
-        assert np.max(final_fieldB.dat.data) <= np.max(true_fieldB.dat.data) + tol, \
-            'Application of the DG1 equispaced limiter in the mixed limiter has not prevented overshoots'
-
-        # Check for no new undershoots in B
-        assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \
-            'Application of the DG1 equispaced limiter in the mixed limiter has not prevented undershoots'
-
-    else:
-        final_field = stepper.fields('tracer')
-
-        # Check tracer is roughly in the correct place
-        assert norm(true_field - final_field) / norm(true_field) < 0.05, \
-            'Something appears to have gone wrong with transport of tracer using a limiter'
+    tol = 1e-9
 
-        # Check for no new overshoots
-        assert np.max(final_field.dat.data) <= np.max(true_field.dat.data) + tol, \
-            'Application of limiter has not prevented overshoots'
+    # Check for no new overshoots
+    assert np.max(final_field.dat.data) <= np.max(true_field.dat.data) + tol, \
+        'Application of limiter has not prevented overshoots'
 
-        # Check for no new undershoots
-        assert np.min(final_field.dat.data) >= np.min(true_field.dat.data) - tol, \
-            'Application of limiter has not prevented undershoots'
+    # Check for no new undershoots
+    assert np.min(final_field.dat.data) >= np.min(true_field.dat.data) - tol, \
+        'Application of limiter has not prevented undershoots'
\ No newline at end of file
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 301c260e5..69c1c55ac 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -62,9 +62,9 @@ def setup_limiters(dirname, space_A, space_B):
     elif space_A == 'HDiv':
         VA = domain.spaces('HDiv')
         space_A_string = 'HDiv'
-    #elif space_A == 'CG1'
-        #VA = FunctionSpace(mesh, 'CG', 1)
-        #space_A_string = 'CG'
+    elif space_A == 'CG1':
+        VA = FunctionSpace(mesh, 'CG', 1)
+        space_A_string = 'CG'
     else:
         raise NotImplementedError
     
@@ -140,6 +140,9 @@ def setup_limiters(dirname, space_A, space_B):
     elif space_A == 'HDiv':
         ibp_A = IntegrateByParts.TWICE
         suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)})
+    elif space_A == 'CG1':
+        ibp_A = IntegrateByParts.NEVER
+        suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)})
     else:
         raise NotImplementedError
         
@@ -182,7 +185,8 @@ def setup_limiters(dirname, space_A, space_B):
     #transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
     
     # DG Upwind transport for both tracers:
-    if space_A == 'HDiv':
+    if space_A == 'HDiv' or space_A == 'CG1':
+        #Pass SUPG options to the transport method
         transport_method_A = DGUpwind(eqn, 'tracerA', ibp=ibp_A)
     else: 
         transport_method_A = DGUpwind(eqn, 'tracerA')
@@ -292,9 +296,9 @@ def setup_limiters(dirname, space_A, space_B):
     return stepper, tmax, true_fieldA, true_fieldB
 
 
-#@pytest.mark.parametrize('space_A', ['HDiv'])#, 
-@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-                                     'DG1', 'DG1_equispaced', 'HDiv'])
+@pytest.mark.parametrize('space_A', ['HDiv'])#, 'CG1'])#, 
+#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
+#                                     'DG1', 'DG1_equispaced', 'HDiv'])
 # It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 

From 6d2fc9b5eb7ce57892e858bddfae3f943ff2e13e Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 23:37:56 +0000
Subject: [PATCH 12/38] finalising mixed options test script

---
 gusto/time_discretisation.py                  | 64 ++-----------------
 gusto/wrappers.py                             | 54 ++++++----------
 .../transport/test_mixed_fs_options.py        | 41 +++---------
 3 files changed, 34 insertions(+), 125 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 0ce9dc68e..2b8f7bb57 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -99,7 +99,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                     elif suboption.name == "recovered":
                         self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)})
                     elif suboption.name == "supg":
-                        self.wrapper.subwrappers.update({field: SUPGWrapper(self, suboption)})
+                        raise RuntimeError(
+                            f'Time discretisation: suboption SUPG is currently not implemented within MixedOptions')
                     else:
                         raise RuntimeError(
                             f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
@@ -180,11 +181,11 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             if type(self.wrapper) == MixedOptions:
 
                 for field, subwrapper in self.wrapper.subwrappers.items():
-                    field_idx = equation.field_names.index(field)
 
-                    subwrapper.idx = field_idx
                     subwrapper.mixed_options = True
 
+                    field_idx = equation.field_names.index(field)
+
                     # Store the original space of the tracer
                     subwrapper.tracer_fs = self.equation.spaces[field_idx]
 
@@ -193,67 +194,16 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
 
-                    # Store test space?
-                    #self.wrapper.test_spaces[field_idx] = subwrapper.function_space
-                    
-                    # Replace the test function space
-                    #if self.wrapper.suboptions[field].name == "supg":
-                    #    new_test = subwrapper.test
-                    #else:
-                    #    new_test = TestFunction(subwrapper.test_space)
-                        
-                    #self.residual = self.residual.label_map(
-                    #    lambda t: t.has_label(transport) and t.get(prognostic) == field,
-                    #    map_if_true=replace_test_function(new_test, old_idx=field_idx))
-                        
-                    #self.residual = subwrapper.label_terms(self.residual)
-                
                 self.wrapper.setup()
-                    
                 self.fs = self.wrapper.function_space
-                
                 new_test_mixed = TestFunctions(self.fs)
 
-                # If use supg, then change the required test function
-
-                # Replace one-by-one in this case?
-                print(type(new_test_mixed))
-
-                for field, subwrapper in self.wrapper.subwrappers.items():
-                    if self.wrapper.suboptions[field].name == "supg":
-                        field_idx = equation.field_names.index(field)
-                        test_list = list(new_test_mixed)
-                        test_list[field_idx] = subwrapper.test
-                        new_test_mixed2 = tuple(test_list)
-                        
-                print(type(new_test_mixed2))
-                        
-                if new_test_mixed == new_test_mixed2:
-                    print('same')
-                else:
-                    print('different')
-                        
-                # Change test functions with the new space:
-                #for field, subwrapper in self.wrapper.subwrappers.items():        
-                #    self.residual = self.residual.label_map(
-                #        lambda t: t.has_label(transport) and t.get(prognostic) == field,
-                #        map_if_true=replace_test_function(new_test_mixed[field_idx], old_idx=field_idx))
-                        
-                #    self.residual = subwrapper.label_terms(self.residual)
-                    
-                # Or, all-at-once:
-                #self.residual = self.residual.label_map(
-                #        lambda t: t.has_label(transport),
-                #        map_if_true=replace_test_function(new_test_mixed))
-                        
+                # Replace the original test function with one from the new
+                # function space defined by the subwrappers
                 self.residual = self.residual.label_map(
                         all_terms,
                         map_if_true=replace_test_function(new_test_mixed))
-                        
-                
-                    
-                    
-                
+
             else:
                 self.wrapper.setup()
                 self.fs = self.wrapper.function_space
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 900224733..880630cd4 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -295,25 +295,9 @@ def setup(self):
             'SUPG wrapper can only be used with SUPG Options'
 
         domain = self.time_discretisation.domain
-        
-        #if self.mixed_options == True:
-        #    self.function_space = self.tracer_fs
-        #else:
-        #    self.function_space = self.time_discretisation.fs
-            
-        # But this will create some problems when defining the new mixed space...
         self.function_space = self.time_discretisation.fs
-        #self.function_space = self.tracer_fs
-        
-        if self.mixed_options == True:
-            self.test_space = self.tracer_fs
-            self.x_out = Function(self.tracer_fs)
-        else:
-            self.test_space = self.function_space
-            self.x_out = Function(self.function_space)
-        
-        #self.test_space = self.function_space
-        #self.x_out = Function(self.function_space)
+        self.test_space = self.function_space
+        self.x_out = Function(self.function_space)
 
         # -------------------------------------------------------------------- #
         # Work out SUPG parameter
@@ -332,8 +316,11 @@ def setup(self):
             # so that we don't apply supg in that direction
             if is_cg(self.function_space):
                 vals = default_vals
+                print('is cg')
             else:
+                print('is not cg')
                 space = self.function_space.ufl_element().sobolev_space
+                print(space.name)
                 if space.name in ["HDiv", "DirectionalH"]:
                     vals = [default_vals[i] if space[i].name == "H1"
                             else 0. for i in range(dim)]
@@ -415,12 +402,11 @@ def __init__(self, equation, suboptions):
         """
         Args:
             equation (:class: `PrognosticEquationSet`): the prognostic equation(s)
-            sublimiters (dict): A dictionary holding limiters defined for individual prognostic variables
+            suboptions (dict): A dictionary holding options defined for individual prognostic variables
         Raises:
             ValueError: If an option is defined for a field that is not in the prognostic variable set
         """
         self.wrapper_spaces = equation.spaces
-        #self.test_spaces = equation.spaces
         self.field_names = equation.field_names
         self.suboptions = suboptions
         self.subwrappers = {}
@@ -431,49 +417,45 @@ def __init__(self, equation, suboptions):
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
     
     def setup(self):
-        
-        # Compute the new mixed function space
+        """ Compute the new mixed function space from the subwrappers """
+
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
-        self.x_out = Function(self.function_space) 
-        #self.test_space = MixedFunctionSpace(self.test_spaces)
-    
+        self.x_out = Function(self.function_space)
+
     def pre_apply(self, x_in):
         """
         Perform the pre-applications for all fields
         with an associated subwrapper.
         """
-            
+
         for field_name in self.field_names:
-        
             field_idx = self.field_names.index(field_name)
             field = x_in.subfunctions[field_idx]
             x_in_sub = self.x_in.subfunctions[field_idx]
-            
+
             if field_name in self.subwrappers:
                 subwrapper = self.subwrappers[field_name]
                 subwrapper.pre_apply(field)
                 x_in_sub.assign(subwrapper.x_in)
             else:
-                x_in_sub.assign(field)        
-            
-            
+                x_in_sub.assign(field)
+
+
     def post_apply(self, x_out):
         """
         Perform the post-applications for all fields
         with an associated subwrapper.
         """
-            
+
         for field_name in self.field_names:
-        
             field_idx = self.field_names.index(field_name)
-            
             field = self.x_out.subfunctions[field_idx]
             x_out_sub = x_out.subfunctions[field_idx]
-            
+
             if field_name in self.subwrappers:
                 subwrapper = self.subwrappers[field_name]
                 subwrapper.x_out.assign(field)
                 subwrapper.post_apply(x_out_sub)
             else:
-                x_out_sub.assign(field)  
\ No newline at end of file
+                x_out_sub.assign(field)
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 69c1c55ac..d00195db2 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -32,7 +32,7 @@ def setup_limiters(dirname, space_A, space_B):
     # Domain
     m = PeriodicIntervalMesh(20, Ld)
     mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20))
-    degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] else 1
+    degree = 0 if space_A in ['DG0', 'Vtheta_degree_0', 'Hdiv'] else 1
 
     domain = Domain(mesh, dt, family="CG", degree=degree)
     
@@ -59,12 +59,6 @@ def setup_limiters(dirname, space_A, space_B):
     elif space_A == 'Vtheta_degree_1':
         VA = domain.spaces('theta')
         space_A_string = 'theta'
-    elif space_A == 'HDiv':
-        VA = domain.spaces('HDiv')
-        space_A_string = 'HDiv'
-    elif space_A == 'CG1':
-        VA = FunctionSpace(mesh, 'CG', 1)
-        space_A_string = 'CG'
     else:
         raise NotImplementedError
     
@@ -137,12 +131,6 @@ def setup_limiters(dirname, space_A, space_B):
     elif space_A == 'Vtheta_degree_1':
         suboptions.update({'tracerA': EmbeddedDGOptions()})
         sublimiters.update({'tracerA': ThetaLimiter(VA)})
-    elif space_A == 'HDiv':
-        ibp_A = IntegrateByParts.TWICE
-        suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)})
-    elif space_A == 'CG1':
-        ibp_A = IntegrateByParts.NEVER
-        suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)})
     else:
         raise NotImplementedError
         
@@ -176,27 +164,16 @@ def setup_limiters(dirname, space_A, space_B):
     else:
         raise NotImplementedError
 
+    # Create the mixed options and mixed limiter objects
     opts = MixedOptions(eqn, suboptions)
     MixedLimiter = MixedFSLimiter(eqn, sublimiters)
 
     # Give the scheme for the coupled transport
-    #transport_schemes = SSPRK3(domain, options=opts)
     transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter)
-    #transport_schemes = SSPRK3(domain, limiter=MixedLimiter)
-    
-    # DG Upwind transport for both tracers:
-    if space_A == 'HDiv' or space_A == 'CG1':
-        #Pass SUPG options to the transport method
-        transport_method_A = DGUpwind(eqn, 'tracerA', ibp=ibp_A)
-    else: 
-        transport_method_A = DGUpwind(eqn, 'tracerA')
-    transport_method_B = DGUpwind(eqn, 'tracerB')
-    
-    transport_method = [transport_method_A, transport_method_B]
-    
-    # Need to give SUPG options to the above, if using supg ...
-    # Need to test SUPG here!
-    
+
+    # DG Upwind transport for both tracers:    
+    transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
+
     # Build time stepper
     stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method)
 
@@ -296,9 +273,9 @@ def setup_limiters(dirname, space_A, space_B):
     return stepper, tmax, true_fieldA, true_fieldB
 
 
-@pytest.mark.parametrize('space_A', ['HDiv'])#, 'CG1'])#, 
-#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
-#                                     'DG1', 'DG1_equispaced', 'HDiv'])
+
+@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
+                                     'DG1', 'DG1_equispaced'])
 # It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
 

From 25ed24320a64bf100814a9531189d010cf398616 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 23:44:30 +0000
Subject: [PATCH 13/38] lint

---
 gusto/time_discretisation.py                        | 13 +++++--------
 gusto/wrappers.py                                   | 10 +++++-----
 .../transport/test_mixed_fs_options.py              |  2 --
 3 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 2b8f7bb57..10ddc8526 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -91,16 +91,13 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                 self.wrapper = options
 
                 for field, suboption in self.wrapper.suboptions.items():
-                    print(field)
-                    print(suboption)
-
                     if suboption.name == 'embedded_dg':
                         self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)})
                     elif suboption.name == "recovered":
                         self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)})
                     elif suboption.name == "supg":
                         raise RuntimeError(
-                            f'Time discretisation: suboption SUPG is currently not implemented within MixedOptions')
+                            'Time discretisation: suboption SUPG is currently not implemented within MixedOptions')
                     else:
                         raise RuntimeError(
                             f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
@@ -201,8 +198,8 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 # Replace the original test function with one from the new
                 # function space defined by the subwrappers
                 self.residual = self.residual.label_map(
-                        all_terms,
-                        map_if_true=replace_test_function(new_test_mixed))
+                    all_terms,
+                    map_if_true=replace_test_function(new_test_mixed))
 
             else:
                 self.wrapper.setup()
@@ -213,12 +210,12 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 # SUPG has a special wrapper
                 if self.wrapper_name == "supg":
                     new_test = self.wrapper.test
-    
+
                 # Replace the original test function with the one from the wrapper
                 self.residual = self.residual.label_map(
                     all_terms,
                     map_if_true=replace_test_function(new_test))
-    
+
                 self.residual = self.wrapper.label_terms(self.residual)
 
         # -------------------------------------------------------------------- #
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 880630cd4..dc2923e9d 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -84,10 +84,10 @@ def setup(self):
 
         assert isinstance(self.options, EmbeddedDGOptions), \
             'Embedded DG wrapper can only be used with Embedded DG Options'
-        
+
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
-        
+
         if self.mixed_options == True:
             original_space = self.tracer_fs
         else:
@@ -111,8 +111,8 @@ def setup(self):
 
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
-        
-        if self.mixed_options == True:
+
+        if self.mixed_options:
             self.x_projected = Function(self.tracer_fs)
         elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
@@ -178,7 +178,7 @@ def setup(self):
 
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
-        
+
         if self.mixed_options == True:
             original_space = self.tracer_fs
         else:
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index d00195db2..ac27f4bd7 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -317,5 +317,3 @@ def test_mixed_fs_options(tmpdir, space_A, space_B):
     # Check for no new undershoots in B
     assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \
         f"Application of the {space_B} equispaced limiter in the mixed limiter has not prevented undershoots"
-
-    
\ No newline at end of file

From f410d2cc01efa49be4f08ccd90d3d14352d351d4 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 23:46:44 +0000
Subject: [PATCH 14/38] lint

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

diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index dc2923e9d..4d42d1dac 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -88,7 +88,7 @@ def setup(self):
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if self.mixed_options == True:
+        if self.mixed_options:
             original_space = self.tracer_fs
         else:
             original_space = self.time_discretisation.fs
@@ -179,7 +179,7 @@ def setup(self):
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if self.mixed_options == True:
+        if self.mixed_options:
             original_space = self.tracer_fs
         else:
             original_space = self.time_discretisation.fs
@@ -201,15 +201,15 @@ def setup(self):
         # Internal variables to be used
         # -------------------------------------------------------------------- #
 
-        if self.mixed_options == True:
+        if self.mixed_options:
             self.x_in_tmp = Function(self.tracer_fs)
         else:
             self.x_in_tmp = Function(self.time_discretisation.fs)
-            
+
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
-        
-        if self.mixed_options == True:
+
+        if self.mixed_options:
             self.x_projected = Function(self.tracer_fs)
         elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
@@ -388,16 +388,16 @@ def label_terms(self, residual):
         new_residual = transporting_velocity.update_value(new_residual, self.transporting_velocity)
 
         return new_residual
-        
+
 
 class MixedOptions(object):
     """
     An object to hold a dictionary with different options for different
     tracers. This means that different tracers can be solved simultaneously
     using a CoupledTransportEquation, whilst being in different spaces
-    and needing different implementation options. 
+    and needing different implementation options.
     """
-    
+
     def __init__(self, equation, suboptions):
         """
         Args:
@@ -410,12 +410,12 @@ def __init__(self, equation, suboptions):
         self.field_names = equation.field_names
         self.suboptions = suboptions
         self.subwrappers = {}
-        
+
         for field, suboption in suboptions.items():
                 # Check that the field is in the prognostic variable set:
                 if field not in equation.field_names:
                     raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
-    
+
     def setup(self):
         """ Compute the new mixed function space from the subwrappers """
 

From d166a224bdfa40d8ffd988823fa9b80565601c2c Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 23:48:35 +0000
Subject: [PATCH 15/38] lint

---
 gusto/wrappers.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 4d42d1dac..7cbcc5af7 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -412,9 +412,9 @@ def __init__(self, equation, suboptions):
         self.subwrappers = {}
 
         for field, suboption in suboptions.items():
-                # Check that the field is in the prognostic variable set:
-                if field not in equation.field_names:
-                    raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
+            # Check that the field is in the prognostic variable set:
+            if field not in equation.field_names:
+                raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
 
     def setup(self):
         """ Compute the new mixed function space from the subwrappers """
@@ -441,7 +441,6 @@ def pre_apply(self, x_in):
             else:
                 x_in_sub.assign(field)
 
-
     def post_apply(self, x_out):
         """
         Perform the post-applications for all fields

From aa3d7c07b25ec7022ac5c6d78b8eea52f25e4822 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 23:53:44 +0000
Subject: [PATCH 16/38] lint

---
 integration-tests/transport/test_limiters.py  |  2 +-
 .../transport/test_mixed_fs_options.py        | 53 ++++++++-----------
 2 files changed, 22 insertions(+), 33 deletions(-)

diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py
index 94d83ae92..8255ad2cd 100644
--- a/integration-tests/transport/test_limiters.py
+++ b/integration-tests/transport/test_limiters.py
@@ -206,4 +206,4 @@ def test_limiters(tmpdir, space):
 
     # Check for no new undershoots
     assert np.min(final_field.dat.data) >= np.min(true_field.dat.data) - tol, \
-        'Application of limiter has not prevented undershoots'
\ No newline at end of file
+        'Application of limiter has not prevented undershoots'
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index ac27f4bd7..96b28c43b 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -35,10 +35,10 @@ def setup_limiters(dirname, space_A, space_B):
     degree = 0 if space_A in ['DG0', 'Vtheta_degree_0', 'Hdiv'] else 1
 
     domain = Domain(mesh, dt, family="CG", degree=degree)
-    
+
     # Transporting velocity space
     V = domain.spaces('HDiv')
-    
+
     # Tracer A spaces
     if space_A == 'DG0':
         VA = domain.spaces('DG')
@@ -61,7 +61,7 @@ def setup_limiters(dirname, space_A, space_B):
         space_A_string = 'theta'
     else:
         raise NotImplementedError
-    
+
     # Tracer B spaces
     if space_B == 'DG':
         space_B_string = 'DG'
@@ -75,7 +75,7 @@ def setup_limiters(dirname, space_A, space_B):
             raise NotImplementedError
     elif space_B == 'Vtheta':
         space_B_string = 'theta'
-        if degree == 0: 
+        if degree == 0:
             VB = domain.spaces('theta')
             VCG1_B = FunctionSpace(mesh, 'CG', 1)
             VDG1_B = domain.spaces('DG1_equispaced')
@@ -85,23 +85,22 @@ def setup_limiters(dirname, space_A, space_B):
             raise NotImplementedError
     else:
         raise NotImplementedError
-        
-    Vpsi = domain.spaces('H1')     
-        
+
+    Vpsi = domain.spaces('H1')
+
     tracerA = ActiveTracer(name='tracerA', space=space_A_string,
-                               variable_type=TracerVariableType.mixing_ratio,
-                               transport_eqn=TransportEquationType.advective)
-    
+                           variable_type=TracerVariableType.mixing_ratio,
+                           transport_eqn=TransportEquationType.advective)
+
     tracerB = ActiveTracer(name='tracerB', space=space_B_string,
-                               variable_type=TracerVariableType.mixing_ratio,
-                               transport_eqn=TransportEquationType.advective)
-        
-        
+                           variable_type=TracerVariableType.mixing_ratio,
+                           transport_eqn=TransportEquationType.advective)
+
     tracers = [tracerA, tracerB]
 
     eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V)
     output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1,
-                                  dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB'])
+                              dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB'])
     
     io = IO(domain, output)
 
@@ -113,21 +112,17 @@ def setup_limiters(dirname, space_A, space_B):
     sublimiters = {}
 
     # Options and limiters for tracer_A
-
     if space_A in ['DG0', 'Vtheta_degree_0']:
         suboptions.update({'tracerA': RecoveryOptions(embedding_space=VDG1_A,
-                                              recovered_space=VCG1_A,
-                                              project_low_method='recover',
-                                              boundary_method=BoundaryMethod.taylor)})
-                                   
+                                                      recovered_space=VCG1_A,
+                                                      project_low_method='recover',
+                                                      boundary_method=BoundaryMethod.taylor)})
+
         sublimiters.update({'tracerA': VertexBasedLimiter(VDG1_A)})
-        
     elif space_A == 'DG1':
         sublimiters.update({'tracerA': DG1Limiter(VA)})
-
     elif space_A == 'DG1_equispaced':
         sublimiters.update({'tracerA': VertexBasedLimiter(VA)})
-
     elif space_A == 'Vtheta_degree_1':
         suboptions.update({'tracerA': EmbeddedDGOptions()})
         sublimiters.update({'tracerA': ThetaLimiter(VA)})
@@ -135,14 +130,12 @@ def setup_limiters(dirname, space_A, space_B):
         raise NotImplementedError
         
     # Options and limiters for tracer_B
-    
     if space_B == 'DG':
         if degree == 0:
             suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B,
-                                              recovered_space=VCG1_B,
-                                              project_low_method='recover',
-                                              boundary_method=BoundaryMethod.taylor)})
-                                   
+                                                          recovered_space=VCG1_B,
+                                                          project_low_method='recover',
+                                                          boundary_method=BoundaryMethod.taylor)})
             sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
         elif degree == 1:
             sublimiters.update({'tracerB': DG1Limiter(VB)})
@@ -273,13 +266,9 @@ def setup_limiters(dirname, space_A, space_B):
     return stepper, tmax, true_fieldA, true_fieldB
 
 
-
 @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0',
                                      'DG1', 'DG1_equispaced'])
-# It only makes sense to use the same degree for tracer B
 @pytest.mark.parametrize('space_B', ['Vtheta', 'DG'])
-
-
 def test_mixed_fs_options(tmpdir, space_A, space_B):
 
     # Setup and run

From f1e958a8101e5b8fa2805d55d6df0dc98b7f84bf Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Mon, 22 Jan 2024 23:55:18 +0000
Subject: [PATCH 17/38] lint

---
 .../transport/test_mixed_fs_options.py          | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 96b28c43b..1c4f545d2 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -101,7 +101,7 @@ def setup_limiters(dirname, space_A, space_B):
     eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V)
     output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1,
                               dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB'])
-    
+
     io = IO(domain, output)
 
     # ------------------------------------------------------------------------ #
@@ -128,7 +128,7 @@ def setup_limiters(dirname, space_A, space_B):
         sublimiters.update({'tracerA': ThetaLimiter(VA)})
     else:
         raise NotImplementedError
-        
+
     # Options and limiters for tracer_B
     if space_B == 'DG':
         if degree == 0:
@@ -142,12 +142,11 @@ def setup_limiters(dirname, space_A, space_B):
         else:
             raise NotImplementedError
     elif space_B == 'Vtheta':
-        if degree == 0: 
+        if degree == 0:
             suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B,
-                                              recovered_space=VCG1_B,
-                                              project_low_method='recover',
-                                              boundary_method=BoundaryMethod.taylor)})
-                                   
+                                                          recovered_space=VCG1_B,
+                                                          project_low_method='recover',
+                                                          boundary_method=BoundaryMethod.taylor)})
             sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)})
         elif degree == 1:
             suboptions.update({'tracerB': EmbeddedDGOptions()})
@@ -164,7 +163,7 @@ def setup_limiters(dirname, space_A, space_B):
     # Give the scheme for the coupled transport
     transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter)
 
-    # DG Upwind transport for both tracers:    
+    # DG Upwind transport for both tracers:
     transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')]
 
     # Build time stepper
@@ -178,7 +177,7 @@ def setup_limiters(dirname, space_A, space_B):
     tracerB_0 = stepper.fields('tracerB', space=VB)
     true_fieldA = stepper.fields('true_tracerA', space=VA)
     true_fieldB = stepper.fields('true_tracerB', space=VB)
-    
+
     x, z = SpatialCoordinate(mesh)
 
     tracer_min = 12.6

From 51f6f7b511c7851c2e7f187d3aba78c81545e511 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 23 Jan 2024 01:29:29 +0000
Subject: [PATCH 18/38] remove dubugging statements

---
 gusto/wrappers.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 7cbcc5af7..110b774df 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -192,7 +192,6 @@ def setup(self):
             V_elt = BrokenElement(original_space.ufl_element())
             self.function_space = FunctionSpace(domain.mesh, V_elt)
         else:
-            print('using embedded space')
             self.function_space = self.options.embedding_space
 
         self.test_space = self.function_space
@@ -316,11 +315,8 @@ def setup(self):
             # so that we don't apply supg in that direction
             if is_cg(self.function_space):
                 vals = default_vals
-                print('is cg')
             else:
-                print('is not cg')
                 space = self.function_space.ufl_element().sobolev_space
-                print(space.name)
                 if space.name in ["HDiv", "DirectionalH"]:
                     vals = [default_vals[i] if space[i].name == "H1"
                             else 0. for i in range(dim)]

From c83e33d8c5a43ede3d09dbcac13e33382da1fc02 Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Wed, 24 Jan 2024 12:08:26 +0000
Subject: [PATCH 19/38] linear solver for the moist convective SW equations

---
 gusto/linear_solvers.py | 106 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 105 insertions(+), 1 deletion(-)

diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py
index d1017aa28..4804010c6 100644
--- a/gusto/linear_solvers.py
+++ b/gusto/linear_solvers.py
@@ -24,7 +24,7 @@
 from abc import ABCMeta, abstractmethod, abstractproperty
 
 
-__all__ = ["IncompressibleSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver"]
+__all__ = ["IncompressibleSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver", "MoistConvectiveSWSolver"]
 
 
 class TimesteppingSolver(object, metaclass=ABCMeta):
@@ -762,3 +762,107 @@ def solve(self, xrhs, dy):
         self.xrhs.assign(xrhs)
         self.solver.solve()
         dy.assign(self.dy)
+
+
+class MoistConvectiveSWSolver(TimesteppingSolver):
+    """
+    Linear solver for the moist convective shallow water equations.
+
+    This solves a linear problem for the shallow water equations with prognostic
+    variables u (velocity) and D (depth). The linear system is solved using a
+    hybridised-mixed method.
+    """
+
+    solver_parameters = {
+        'ksp_type': 'preonly',
+        'mat_type': 'matfree',
+        'pc_type': 'python',
+        'pc_python_type': 'firedrake.HybridizationPC',
+        'hybridization': {'ksp_type': 'cg',
+                          'pc_type': 'gamg',
+                          'ksp_rtol': 1e-8,
+                          'mg_levels': {'ksp_type': 'chebyshev',
+                                        'ksp_max_it': 2,
+                                        'pc_type': 'bjacobi',
+                                        'sub_pc_type': 'ilu'}}
+    }
+
+    @timed_function("Gusto:SolverSetup")
+    def _setup_solver(self):
+        equation = self.equations      # just cutting down line length a bit
+        dt = self.dt
+        beta_ = dt*self.alpha
+        Vu = equation.domain.spaces("HDiv")
+        VD = equation.domain.spaces("DG")
+
+        # Store time-stepping coefficients as UFL Constants
+        beta = Constant(beta_)
+
+        # Split up the rhs vector
+        self.xrhs = Function(self.equations.function_space)
+        u_in = split(self.xrhs)[0]
+        D_in = split(self.xrhs)[1]
+
+        # Build the reduced function space for u, D
+        M = MixedFunctionSpace((Vu, VD))
+        w, phi = TestFunctions(M)
+        u, D = TrialFunctions(M)
+
+        # Get background depth
+        Dbar = split(equation.X_ref)[1]
+
+        g = equation.parameters.g
+
+        eqn = (
+            inner(w, (u - u_in)) * dx
+            - beta * (D - Dbar) * div(w*g) * dx
+            + inner(phi, (D - D_in)) * dx
+            + beta * phi * Dbar * div(u) * dx
+        )
+
+        aeqn = lhs(eqn)
+        Leqn = rhs(eqn)
+
+        # Place to put results of (u,D) solver
+        self.uD = Function(M)
+
+        # Boundary conditions
+        bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']]
+
+        # Solver for u, D
+        uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs)
+
+        # Provide callback for the nullspace of the trace system
+        def trace_nullsp(T):
+            return VectorSpaceBasis(constant=True)
+
+        appctx = {"trace_nullspace": trace_nullsp}
+        self.uD_solver = LinearVariationalSolver(uD_problem,
+                                                 solver_parameters=self.solver_parameters,
+                                                 appctx=appctx)
+
+        # Log residuals on hybridized solver
+        self.log_ksp_residuals(self.uD_solver.snes.ksp)
+
+    @timed_function("Gusto:LinearSolve")
+    def solve(self, xrhs, dy):
+        """
+        Solve the linear problem.
+
+        Args:
+            xrhs (:class:`Function`): the right-hand side field in the
+                appropriate :class:`MixedFunctionSpace`.
+            dy (:class:`Function`): the resulting field in the appropriate
+                :class:`MixedFunctionSpace`.
+        """
+        self.xrhs.assign(xrhs)
+
+        with timed_region("Gusto:VelocityDepthSolve"):
+            logger.info('Moist convective linear solver: mixed solve')
+            self.uD_solver.solve()
+
+        u1, D1 = self.uD.subfunctions
+        u = dy.subfunctions[0]
+        D = dy.subfunctions[1]
+        u.assign(u1)
+        D.assign(D1)

From 50fbb956ad70e693c10ab93f094240c42924ddb8 Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Fri, 16 Feb 2024 11:51:13 +0000
Subject: [PATCH 20/38] include missing spatial discretisation in moist thermal
 example

---
 examples/shallow_water/moist_thermal_williamson5.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py
index c0626063c..9fcf2423f 100644
--- a/examples/shallow_water/moist_thermal_williamson5.py
+++ b/examples/shallow_water/moist_thermal_williamson5.py
@@ -104,8 +104,10 @@ def gamma_v(x_in):
 InstantRain(eqns, qprecip, vapour_name="cloud_water", rain_name="rain",
             gamma_r=gamma_r)
 
+transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names]
+
 # Timestepper
-stepper = Timestepper(eqns, RK4(domain), io)
+stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods)
 
 # ----------------------------------------------------------------- #
 # Initial conditions

From 850f6cce959e3d8f69c15812f86682e62a00f7af Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Mon, 19 Feb 2024 14:20:03 +0000
Subject: [PATCH 21/38] moist convective version of Williamson 2, using the
 SIQN stepper and the moist convective linear solver

---
 .../moist_convective_williamson2.py           | 164 ++++++++++++++++++
 1 file changed, 164 insertions(+)
 create mode 100644 examples/shallow_water/moist_convective_williamson2.py

diff --git a/examples/shallow_water/moist_convective_williamson2.py b/examples/shallow_water/moist_convective_williamson2.py
new file mode 100644
index 000000000..f02e992dd
--- /dev/null
+++ b/examples/shallow_water/moist_convective_williamson2.py
@@ -0,0 +1,164 @@
+"""
+A moist convective version of the Williamson 2 shallow water test (steady state
+geostrophically-balanced flow). The saturation function depends on height,
+with a constant background buoyancy/temperature field.
+Vapour is initialised very close to saturation and small overshoots will
+generate clouds.
+"""
+from gusto import *
+from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, sin, cos, exp)
+import sys
+
+# ----------------------------------------------------------------- #
+# Test case parameters
+# ----------------------------------------------------------------- #
+
+dt = 120
+
+if '--running-tests' in sys.argv:
+    tmax = dt
+    dumpfreq = 1
+else:
+    day = 24*60*60
+    tmax = 5*day
+    ndumps = 5
+    dumpfreq = int(tmax / (ndumps*dt))
+
+R = 6371220.
+u_max = 20
+phi_0 = 3e4
+epsilon = 1/300
+theta_0 = epsilon*phi_0**2
+g = 9.80616
+H = phi_0/g
+xi = 0
+q0 = 200
+beta1 = 110
+alpha = 16
+gamma_v = 0.98
+qprecip = 1e-4
+gamma_r = 1e-3
+
+# ----------------------------------------------------------------- #
+# Set up model objects
+# ----------------------------------------------------------------- #
+
+# Domain
+mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2)
+degree = 1
+domain = Domain(mesh, dt, 'BDM', degree)
+x = SpatialCoordinate(mesh)
+
+# Equations
+parameters = ShallowWaterParameters(H=H, g=g)
+Omega = parameters.Omega
+fexpr = 2*Omega*x[2]/R
+
+tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')]
+
+eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr,
+                             u_transport_option='vector_advection_form',
+                             active_tracers=tracers)
+
+# IO
+dirname = "moist_convective_williamson2"
+dumpfreq = int(tmax / (ndumps*dt))
+output = OutputParameters(dirname=dirname,
+                          dumpfreq=dumpfreq,
+                          dumplist_latlon=['D', 'D_error'],
+                          dump_nc=True,
+                          dump_vtus=True)
+
+diagnostic_fields = [CourantNumber(), RelativeVorticity(),
+                     PotentialVorticity(),
+                     ShallowWaterKineticEnergy(),
+                     ShallowWaterPotentialEnergy(parameters),
+                     ShallowWaterPotentialEnstrophy(),
+                     SteadyStateError('u'), SteadyStateError('D'),
+                     SteadyStateError('water_vapour'),
+                     SteadyStateError('cloud_water')]
+
+io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+
+
+# define saturation function
+def sat_func(x_in):
+    h = x_in.split()[1]
+    lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2])
+    numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma))
+    denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2
+    theta = numerator/denominator
+    return q0/(g*h) * exp(20*(theta))
+
+
+transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names]
+
+limiter = DG1Limiter(domain.spaces('DG'))
+
+transported_fields = [TrapeziumRule(domain, "u"),
+                      SSPRK3(domain, "D"),
+                      SSPRK3(domain, "water_vapour", limiter=limiter),
+                      SSPRK3(domain, "cloud_water", limiter=limiter),
+                      SSPRK3(domain, "rain", limiter=limiter)
+                      ]
+
+linear_solver = MoistConvectiveSWSolver(eqns)
+
+sat_adj = SWSaturationAdjustment(eqns, sat_func,
+                                 time_varying_saturation=True,
+                                 convective_feedback=True, beta1=beta1,
+                                 gamma_v=gamma_v, time_varying_gamma_v=False,
+                                 parameters=parameters)
+
+inst_rain = InstantRain(eqns, qprecip, vapour_name="cloud_water",
+                        rain_name="rain", gamma_r=gamma_r)
+
+physics_schemes = [(sat_adj, ForwardEuler(domain)),
+                   (inst_rain, ForwardEuler(domain))]
+
+stepper = SemiImplicitQuasiNewton(eqns, io,
+                                  transport_schemes=transported_fields,
+                                  spatial_methods=transport_methods,
+                                  linear_solver=linear_solver,
+                                  physics_schemes=physics_schemes)
+
+# ----------------------------------------------------------------- #
+# Initial conditions
+# ----------------------------------------------------------------- #
+
+u0 = stepper.fields("u")
+D0 = stepper.fields("D")
+v0 = stepper.fields("water_vapour")
+
+lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2])
+
+uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x)
+g = parameters.g
+w = Omega*R*u_max + (u_max**2)/2
+sigma = 0
+
+Dexpr = H - (1/g)*(w)*((sin(phi))**2)
+D_for_v = H - (1/g)*(w + sigma)*((sin(phi))**2)
+
+# though this set-up has no buoyancy, we use the expression for theta to set up
+# the initial vapour
+numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma))
+denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2
+theta = numerator/denominator
+
+initial_msat = q0/(g*Dexpr) * exp(20*theta)
+vexpr = (1 - xi) * initial_msat
+
+u0.project(uexpr)
+D0.interpolate(Dexpr)
+v0.interpolate(vexpr)
+
+# Set reference profiles
+Dbar = Function(D0.function_space()).assign(H)
+stepper.set_reference_profiles([('D', Dbar)])
+
+# ----------------------------------------------------------------- #
+# Run
+# ----------------------------------------------------------------- #
+
+stepper.run(t=0, tmax=tmax)

From a1b39377d4b1add2c39f578dd001a7791b959431 Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Mon, 19 Feb 2024 14:38:54 +0000
Subject: [PATCH 22/38] check that the buoyancy field does exist when using the
 thermal linear solver, and add a warning if no reference profile has been set
 for buoyancy in the stepper

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

diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py
index d1017aa28..b26a3bbba 100644
--- a/gusto/linear_solvers.py
+++ b/gusto/linear_solvers.py
@@ -586,6 +586,10 @@ def _setup_solver(self):
         VD = equation.domain.spaces("DG")
         Vb = equation.domain.spaces("DG")
 
+        # Check that the third field is buoyancy
+        if not equation.field_names[2] == 'b':
+            raise NotImplementedError("Field 'b' must exist to use the thermal linear solver in the SIQN scheme")
+
         # Store time-stepping coefficients as UFL Constants
         beta = Constant(beta_)
 
@@ -667,6 +671,13 @@ def solve(self, xrhs, dy):
         """
         self.xrhs.assign(xrhs)
 
+        # Check that the b reference profile has been set
+        bbar = split(self.equations.X_ref)[2]
+        b = dy.subfunctions[2]
+        bbar_func = Function(b.function_space()).interpolate(bbar)
+        if bbar_func.dat.data.max() == 0 and bbar_func.dat.data.min() == 0:
+            logger.warning("The reference profile for b in the linear solver is zero. To set a non-zero profile add b to the set_reference_profiles argument.")
+
         with timed_region("Gusto:VelocityDepthSolve"):
             logger.info('Thermal linear solver: mixed solve')
             self.uD_solver.solve()

From 51b257b087cfdd083c597582b37417d06f9b91c7 Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Mon, 19 Feb 2024 14:50:32 +0000
Subject: [PATCH 23/38] interpolate bbar to the continuous space

---
 gusto/linear_solvers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py
index b26a3bbba..2679593b1 100644
--- a/gusto/linear_solvers.py
+++ b/gusto/linear_solvers.py
@@ -605,7 +605,9 @@ def _setup_solver(self):
         u, D = TrialFunctions(M)
 
         # Get background buoyancy and depth
-        bbar = split(equation.X_ref)[2]
+        _bbar = split(equation.X_ref)[2]
+        VH1 = equation.domain.spaces("H1")
+        bbar = Function(VH1).interpolate(_bbar)
         Dbar = split(equation.X_ref)[1]
 
         # Approximate elimination of b

From f7b8e9f508ac6e7998c9067728cfecddd22fcb85 Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Mon, 19 Feb 2024 17:36:57 +0000
Subject: [PATCH 24/38] remove erroneous repeated definition of dumpfreq

---
 examples/shallow_water/moist_convective_williamson2.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/shallow_water/moist_convective_williamson2.py b/examples/shallow_water/moist_convective_williamson2.py
index f02e992dd..22ca6b78e 100644
--- a/examples/shallow_water/moist_convective_williamson2.py
+++ b/examples/shallow_water/moist_convective_williamson2.py
@@ -62,7 +62,6 @@
 
 # IO
 dirname = "moist_convective_williamson2"
-dumpfreq = int(tmax / (ndumps*dt))
 output = OutputParameters(dirname=dirname,
                           dumpfreq=dumpfreq,
                           dumplist_latlon=['D', 'D_error'],

From bc22c7d1c5d99a246b070c3452deecdbda7a3854 Mon Sep 17 00:00:00 2001
From: nh491 <nh491@pileus.ex.ac.uk>
Date: Mon, 19 Feb 2024 17:42:40 +0000
Subject: [PATCH 25/38] changing qprecip and qr so that we get rain

---
 examples/shallow_water/moist_thermal_williamson5.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py
index 9fcf2423f..314f7a75e 100644
--- a/examples/shallow_water/moist_thermal_williamson5.py
+++ b/examples/shallow_water/moist_thermal_williamson5.py
@@ -34,8 +34,8 @@
 mu2 = 0.98
 q0 = 135  # chosen to give an initial max vapour of approx 0.02
 beta2 = 10
-qprecip = 10e-4
-gamma_r = 10e-3
+qprecip = 1e-4
+gamma_r = 1e-3
 # topography parameters
 R0 = pi/9.
 R0sq = R0**2

From 55e0856b7027924cec6a1a58f1a946a91cac1c86 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 20 Feb 2024 02:40:47 +0000
Subject: [PATCH 26/38] separate MixedFSOptions and MixedFSWrapper to align
 with exising code

---
 gusto/configuration.py                        | 13 +++++-
 gusto/time_discretisation.py                  | 36 ++++++++-------
 gusto/wrappers.py                             | 44 +++++++------------
 .../transport/test_mixed_fs_options.py        |  3 +-
 4 files changed, 48 insertions(+), 48 deletions(-)

diff --git a/gusto/configuration.py b/gusto/configuration.py
index 3694ae9ea..21e16513f 100644
--- a/gusto/configuration.py
+++ b/gusto/configuration.py
@@ -7,7 +7,7 @@
 __all__ = [
     "IntegrateByParts", "TransportEquationType", "OutputParameters",
     "CompressibleParameters", "ShallowWaterParameters",
-    "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions",
+    "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", "MixedFSOptions",
     "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters"
 ]
 
@@ -172,6 +172,17 @@ class SUPGOptions(WrapperOptions):
     ibp = IntegrateByParts.TWICE
 
 
+class MixedFSOptions(WrapperOptions):
+    """Specifies options for a mixed finite element formulation
+    where different suboptions are applied to different
+    prognostic variables."""
+
+    name = "mixed_options"
+    #tracer_fs = None
+    suboptions = {}
+    #subwrappers = {}
+
+
 class SpongeLayerParameters(Configuration):
     """Specifies parameters describing a 'sponge' (damping) layer."""
 
diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 10ddc8526..2d010ffcb 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -87,10 +87,11 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
         self.courant_max = None
 
         if options is not None:
-            if type(options) == MixedOptions:
-                self.wrapper = options
+            self.wrapper_name = options.name
+            if self.wrapper_name == "mixed_options":
+                self.wrapper = MixedFSWrapper()
 
-                for field, suboption in self.wrapper.suboptions.items():
+                for field, suboption in options.suboptions.items():
                     if suboption.name == 'embedded_dg':
                         self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)})
                     elif suboption.name == "recovered":
@@ -101,18 +102,15 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                     else:
                         raise RuntimeError(
                             f'Time discretisation: suboption wrapper {wrapper_name} not implemented')
-
+            elif self.wrapper_name == "embedded_dg":
+                self.wrapper = EmbeddedDGWrapper(self, options)
+            elif self.wrapper_name == "recovered":
+                self.wrapper = RecoveryWrapper(self, options)
+            elif self.wrapper_name == "supg":
+                self.wrapper = SUPGWrapper(self, options)
             else:
-                self.wrapper_name = options.name
-                if self.wrapper_name == "embedded_dg":
-                    self.wrapper = EmbeddedDGWrapper(self, options)
-                elif self.wrapper_name == "recovered":
-                    self.wrapper = RecoveryWrapper(self, options)
-                elif self.wrapper_name == "supg":
-                    self.wrapper = SUPGWrapper(self, options)
-                else:
-                    raise RuntimeError(
-                        f'Time discretisation: wrapper {self.wrapper_name} not implemented')
+                raise RuntimeError(
+                    f'Time discretisation: wrapper {self.wrapper_name} not implemented')
         else:
             self.wrapper = None
             self.wrapper_name = None
@@ -175,11 +173,15 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         # -------------------------------------------------------------------- #
 
         if self.wrapper is not None:
-            if type(self.wrapper) == MixedOptions:
+            if self.wrapper_name == "mixed_options":
+            
+                self.wrapper.wrapper_spaces = equation.spaces
+                self.wrapper.field_names = equation.field_names
 
                 for field, subwrapper in self.wrapper.subwrappers.items():
-
-                    subwrapper.mixed_options = True
+                
+                    if field not in equation.field_names:
+                        raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
 
                     field_idx = equation.field_names.index(field)
 
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 110b774df..0ac2728c3 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -15,7 +15,7 @@
 from gusto.labels import transporting_velocity
 import ufl
 
-__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedOptions"]
+__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedFSWrapper"]
 
 
 class Wrapper(object, metaclass=ABCMeta):
@@ -33,8 +33,6 @@ def __init__(self, time_discretisation, wrapper_options):
         self.time_discretisation = time_discretisation
         self.options = wrapper_options
         self.solver_parameters = None
-        self.idx = None
-        self.mixed_options = False
         self.tracer_fs = None
 
     @abstractmethod
@@ -88,7 +86,7 @@ def setup(self):
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if self.mixed_options:
+        if self.tracer_fs is not None:
             original_space = self.tracer_fs
         else:
             original_space = self.time_discretisation.fs
@@ -112,7 +110,7 @@ def setup(self):
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
 
-        if self.mixed_options:
+        if self.tracer_fs is not None:
             self.x_projected = Function(self.tracer_fs)
         elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
@@ -179,7 +177,7 @@ def setup(self):
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if self.mixed_options:
+        if self.tracer_fs is not None:
             original_space = self.tracer_fs
         else:
             original_space = self.time_discretisation.fs
@@ -200,7 +198,7 @@ def setup(self):
         # Internal variables to be used
         # -------------------------------------------------------------------- #
 
-        if self.mixed_options:
+        if self.tracer_fs is not None:
             self.x_in_tmp = Function(self.tracer_fs)
         else:
             self.x_in_tmp = Function(self.time_discretisation.fs)
@@ -208,7 +206,7 @@ def setup(self):
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
 
-        if self.mixed_options:
+        if self.tracer_fs is not None:
             self.x_projected = Function(self.tracer_fs)
         elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
@@ -386,31 +384,19 @@ def label_terms(self, residual):
         return new_residual
 
 
-class MixedOptions(object):
+class MixedFSWrapper(object):
     """
-    An object to hold a dictionary with different options for different
-    tracers. This means that different tracers can be solved simultaneously
-    using a CoupledTransportEquation, whilst being in different spaces
-    and needing different implementation options.
+    An object to hold a subwrapper dictionary with different wrappers for
+    different tracers. This means that different tracers can be solved
+    simultaneously using a CoupledTransportEquation, whilst being in
+    different spaces and needing different implementation options.
     """
 
-    def __init__(self, equation, suboptions):
-        """
-        Args:
-            equation (:class: `PrognosticEquationSet`): the prognostic equation(s)
-            suboptions (dict): A dictionary holding options defined for individual prognostic variables
-        Raises:
-            ValueError: If an option is defined for a field that is not in the prognostic variable set
-        """
-        self.wrapper_spaces = equation.spaces
-        self.field_names = equation.field_names
-        self.suboptions = suboptions
-        self.subwrappers = {}
+    def __init__(self):
 
-        for field, suboption in suboptions.items():
-            # Check that the field is in the prognostic variable set:
-            if field not in equation.field_names:
-                raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set")
+        self.wrapper_spaces = None
+        self.field_names = None
+        self.subwrappers = {}
 
     def setup(self):
         """ Compute the new mixed function space from the subwrappers """
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 1c4f545d2..097751212 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -157,7 +157,8 @@ def setup_limiters(dirname, space_A, space_B):
         raise NotImplementedError
 
     # Create the mixed options and mixed limiter objects
-    opts = MixedOptions(eqn, suboptions)
+    #opts = MixedFSOptions(eqn, suboptions)
+    opts = MixedFSOptions(suboptions=suboptions)
     MixedLimiter = MixedFSLimiter(eqn, sublimiters)
 
     # Give the scheme for the coupled transport

From 164041cbd71d1c4811e27b7aff61d417a8c62095 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 20 Feb 2024 02:43:54 +0000
Subject: [PATCH 27/38] lint

---
 gusto/configuration.py                               | 2 --
 integration-tests/transport/test_mixed_fs_options.py | 1 -
 2 files changed, 3 deletions(-)

diff --git a/gusto/configuration.py b/gusto/configuration.py
index 21e16513f..5c909aa59 100644
--- a/gusto/configuration.py
+++ b/gusto/configuration.py
@@ -178,9 +178,7 @@ class MixedFSOptions(WrapperOptions):
     prognostic variables."""
 
     name = "mixed_options"
-    #tracer_fs = None
     suboptions = {}
-    #subwrappers = {}
 
 
 class SpongeLayerParameters(Configuration):
diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py
index 097751212..fc33dae6d 100644
--- a/integration-tests/transport/test_mixed_fs_options.py
+++ b/integration-tests/transport/test_mixed_fs_options.py
@@ -157,7 +157,6 @@ def setup_limiters(dirname, space_A, space_B):
         raise NotImplementedError
 
     # Create the mixed options and mixed limiter objects
-    #opts = MixedFSOptions(eqn, suboptions)
     opts = MixedFSOptions(suboptions=suboptions)
     MixedLimiter = MixedFSLimiter(eqn, sublimiters)
 

From edcc929e9518f996e0f68ce0c09f7440b34723ef Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 20 Feb 2024 02:45:43 +0000
Subject: [PATCH 28/38] lint

---
 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 2d010ffcb..49cc46e1d 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -174,7 +174,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         if self.wrapper is not None:
             if self.wrapper_name == "mixed_options":
-            
+
                 self.wrapper.wrapper_spaces = equation.spaces
                 self.wrapper.field_names = equation.field_names
 

From b317f39024bda5137f3aaa459df1f1728fa3ae1c Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 20 Feb 2024 02:47:24 +0000
Subject: [PATCH 29/38] lint

---
 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 49cc46e1d..5d9297c75 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -179,7 +179,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 self.wrapper.field_names = equation.field_names
 
                 for field, subwrapper in self.wrapper.subwrappers.items():
-                
+
                     if field not in equation.field_names:
                         raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
 

From 8800692ffb9fd2ec287b359994b3851c73249e33 Mon Sep 17 00:00:00 2001
From: nhartney <nh491@exeter.ac.uk>
Date: Thu, 22 Feb 2024 16:37:49 +0000
Subject: [PATCH 30/38] discretise the terms involving bbar as though it is not
 in a continuous space

---
 gusto/linear_solvers.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py
index 2679593b1..375921198 100644
--- a/gusto/linear_solvers.py
+++ b/gusto/linear_solvers.py
@@ -8,7 +8,7 @@
 from firedrake import (
     split, LinearVariationalProblem, Constant, LinearVariationalSolver,
     TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs,
-    rhs, FacetNormal, div, dx, jump, avg, dS_v, dS_h, ds_v, ds_t, ds_b,
+    rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b,
     ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis,
     BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC
 )
@@ -605,20 +605,22 @@ def _setup_solver(self):
         u, D = TrialFunctions(M)
 
         # Get background buoyancy and depth
-        _bbar = split(equation.X_ref)[2]
-        VH1 = equation.domain.spaces("H1")
-        bbar = Function(VH1).interpolate(_bbar)
         Dbar = split(equation.X_ref)[1]
+        bbar = split(equation.X_ref)[2]
 
         # Approximate elimination of b
         b = -dot(u, grad(bbar))*beta + b_in
 
+        n = FacetNormal(equation.domain.mesh)
+
         eqn = (
             inner(w, (u - u_in)) * dx
             - beta * (D - Dbar) * div(w*bbar) * dx
-            + beta * 0.5 * Dbar * inner(w, grad(bbar)) * dx
+            + beta * jump(w*bbar, n) * avg(D-Dbar) * dS
+            - beta * 0.5 * Dbar * bbar * div(w) * dx
             - beta * 0.5 * Dbar * b * div(w) * dx
-            + beta * 0.5 * (D - Dbar) * inner(w, grad(bbar)) * dx
+            - beta * 0.5 * bbar * div(w*(D-Dbar)) * dx
+            + beta * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS
             + inner(phi, (D - D_in)) * dx
             + beta * phi * Dbar * div(u) * dx
         )

From 1def9eb61c8784b97eb4a0aa88c9daf59fa53cd3 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Tue, 27 Feb 2024 21:02:54 +0000
Subject: [PATCH 31/38] add a limiter to clip fields to be non-negative

---
 gusto/kernels.py                              |  39 +++-
 gusto/limiters.py                             |  51 ++++-
 .../transport/test_zero_limiter.py            | 174 ++++++++++++++++++
 .../kernel_tests/test_clip_zero_kernel.py     |  40 ++++
 4 files changed, 301 insertions(+), 3 deletions(-)
 create mode 100644 integration-tests/transport/test_zero_limiter.py
 create mode 100644 unit-tests/kernel_tests/test_clip_zero_kernel.py

diff --git a/gusto/kernels.py b/gusto/kernels.py
index e16996f5d..7445d7a68 100644
--- a/gusto/kernels.py
+++ b/gusto/kernels.py
@@ -10,7 +10,7 @@
 """
 
 from firedrake import dx
-from firedrake.parloops import par_loop, READ, WRITE
+from firedrake.parloops import par_loop, READ, WRITE, RW
 
 
 class LimitMidpoints():
@@ -75,3 +75,40 @@ def apply(self, field_hat, field_DG1, field_old):
                   "field_DG1": (field_DG1, READ),
                   "field_old": (field_old, READ)},
                  is_loopy_kernel=True)
+
+
+class ClipZero():
+    """Clips any negative field values to be zero."""
+
+    def __init__(self, V):
+        """
+        Args:
+            V (:class:`FunctionSpace`): The space of the field to be clipped.
+        """
+        shapes = {'nDOFs': V.finat_element.space_dimension()}
+        domain = "{{[i]: 0 <= i < {nDOFs}}}".format(**shapes)
+
+        instrs = ("""
+                  for i
+                      if field_in[i] < 0.0
+                          field[i] = 0.0
+                      else
+                          field[i] = field_in[i]
+                      end
+                  end
+                  """)
+
+        self._kernel = (domain, instrs)
+
+    def apply(self, field, field_in):
+        """
+        Performs the par loop.
+
+        Args:
+            field (:class:`Function`): The field to be written to.
+            field_in (:class:`Function`): The field to be clipped.
+        """
+        par_loop(self._kernel, dx,
+                 {"field": (field, WRITE),
+                  "field_in": (field_in, READ)},
+                 is_loopy_kernel=True)
diff --git a/gusto/limiters.py b/gusto/limiters.py
index a39c6688d..7366c09e2 100644
--- a/gusto/limiters.py
+++ b/gusto/limiters.py
@@ -8,11 +8,12 @@
 from firedrake import (BrokenElement, Function, FunctionSpace, interval,
                        FiniteElement, TensorProductElement)
 from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter
-from gusto.kernels import LimitMidpoints
+from gusto.kernels import LimitMidpoints, ClipZero
 
 import numpy as np
 
-__all__ = ["DG1Limiter", "ThetaLimiter", "NoLimiter", "MixedFSLimiter"]
+__all__ = ["DG1Limiter", "ThetaLimiter", "NoLimiter", "ZeroLimiter",
+           "MixedFSLimiter"]
 
 
 class DG1Limiter(object):
@@ -163,6 +164,52 @@ def apply(self, field):
         field.interpolate(self.field_hat)
 
 
+class ZeroLimiter(object):
+    """
+    A simple limiter to enforce non-negativity of a field pointwise.
+
+    Negative values are simply clipped to be zero. There is also the option to
+    project the field to another function space to enforce non-negativity there.
+    """
+
+    def __init__(self, space, clipping_space=None):
+        """
+        Args:
+            space (:class:`FunctionSpace`): the space of the incoming field to
+                clip.
+            clipping_space (:class:`FunctionSpace`, optional): the space in
+                which to clip the field. If not specified, the space of the
+                input field is used.
+        """
+
+        self.space = space
+        if clipping_space is not None:
+            self.clipping_space = clipping_space
+            self.map_to_clip = True
+            self.field_to_clip = Function(self.clipping_space)
+        else:
+            self.clipping_space = space
+            self.map_to_clip = False
+
+        self._kernel = ClipZero(self.clipping_space)
+
+    def apply(self, field):
+        """
+        The application of the limiter to the field.
+
+        Args:
+            field (:class:`Function`): the field to apply the limiter to.
+         """
+
+        # Obtain field in clipping space
+        if self.map_to_clip:
+            self.field_to_clip.interpolate(field)
+            self._kernel.apply(self.field_to_clip, self.field_to_clip)
+            field.interpolate(self.field_to_clip)
+        else:
+            self._kernel.apply(field, field)
+
+
 class NoLimiter(object):
     """A blank limiter that does nothing."""
 
diff --git a/integration-tests/transport/test_zero_limiter.py b/integration-tests/transport/test_zero_limiter.py
new file mode 100644
index 000000000..2e61b1f02
--- /dev/null
+++ b/integration-tests/transport/test_zero_limiter.py
@@ -0,0 +1,174 @@
+"""
+This tests the ZeroLimiter, which enforces non-negativity.
+A sharp bubble of warm air is generated in a vertical slice and then transported
+by a prescribed transport scheme. If the limiter is working, the transport
+should have produced no negative values.
+"""
+
+from gusto import *
+from firedrake import (as_vector, PeriodicIntervalMesh, pi, SpatialCoordinate,
+                       ExtrudedMesh, FunctionSpace, Function, norm,
+                       conditional, sqrt)
+import numpy as np
+import pytest
+
+
+def setup_zero_limiter(dirname, clipping_space):
+
+    # ------------------------------------------------------------------------ #
+    # Parameters for test case
+    # ------------------------------------------------------------------------ #
+
+    Ld = 1.
+    tmax = 0.2
+    dt = tmax / 40
+    rotations = 0.25
+
+    # ------------------------------------------------------------------------ #
+    # Build model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    m = PeriodicIntervalMesh(20, Ld)
+    mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20))
+    degree = 1
+
+    domain = Domain(mesh, dt, family="CG", degree=degree)
+
+    DG1 = FunctionSpace(mesh, 'DG', 1)
+    DG1_equispaced = domain.spaces('DG1_equispaced')
+
+    Vpsi = domain.spaces('H1')
+
+    eqn = AdvectionEquation(domain, DG1, 'tracer')
+    output = OutputParameters(dirname=dirname+'/limiters',
+                              dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer'])
+
+    io = IO(domain, output)
+
+    # ------------------------------------------------------------------------ #
+    # Set up transport scheme
+    # ------------------------------------------------------------------------ #
+
+    if clipping_space is None:
+        limiter = ZeroLimiter(DG1)
+    elif clipping_space == 'equispaced':
+        limiter = ZeroLimiter(DG1, clipping_space=DG1_equispaced)
+
+    transport_schemes = SSPRK3(domain, limiter=limiter)
+    transport_method = DGUpwind(eqn, "tracer")
+
+    # Build time stepper
+    stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method)
+
+    # ------------------------------------------------------------------------ #
+    # Initial condition
+    # ------------------------------------------------------------------------ #
+
+    tracer0 = stepper.fields('tracer', DG1)
+    true_field = stepper.fields('true_tracer', space=DG1)
+
+    x, z = SpatialCoordinate(mesh)
+
+    tracer_min = 12.6
+    dtracer = 3.2
+
+    # First time do initial conditions, second time do final conditions
+    for i in range(2):
+
+        if i == 0:
+            x1_lower = 2 * Ld / 5
+            x1_upper = 3 * Ld / 5
+            z1_lower = 6 * Ld / 10
+            z1_upper = 8 * Ld / 10
+            x2_lower = 6 * Ld / 10
+            x2_upper = 8 * Ld / 10
+            z2_lower = 2 * Ld / 5
+            z2_upper = 3 * Ld / 5
+        elif i == 1:
+            # Rotated anti-clockwise by 90 degrees (x -> z, z -> -x)
+            x1_lower = 2 * Ld / 10
+            x1_upper = 4 * Ld / 10
+            z1_lower = 2 * Ld / 5
+            z1_upper = 3 * Ld / 5
+            x2_lower = 2 * Ld / 5
+            x2_upper = 3 * Ld / 5
+            z2_lower = 6 * Ld / 10
+            z2_upper = 8 * Ld / 10
+        else:
+            raise ValueError
+
+        expr_1 = conditional(x > x1_lower,
+                             conditional(x < x1_upper,
+                                         conditional(z > z1_lower,
+                                                     conditional(z < z1_upper, dtracer, 0.0),
+                                                     0.0),
+                                         0.0),
+                             0.0)
+
+        expr_2 = conditional(x > x2_lower,
+                             conditional(x < x2_upper,
+                                         conditional(z > z2_lower,
+                                                     conditional(z < z2_upper, dtracer, 0.0),
+                                                     0.0),
+                                         0.0),
+                             0.0)
+
+        if i == 0:
+            tracer0.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+        elif i == 1:
+            true_field.interpolate(Constant(tracer_min) + expr_1 + expr_2)
+        else:
+            raise ValueError
+
+    # ------------------------------------------------------------------------ #
+    # Velocity profile
+    # ------------------------------------------------------------------------ #
+
+    psi = Function(Vpsi)
+    u = stepper.fields('u')
+
+    # set up solid body rotation for transport
+    # we do this slightly complicated stream function to make the velocity 0 at edges
+    # thus we avoid odd effects at boundaries
+    xc = Ld / 2
+    zc = Ld / 2
+    r = sqrt((x - xc) ** 2 + (z - zc) ** 2)
+    omega = rotations * 2 * pi / tmax
+    r_out = 9 * Ld / 20
+    r_in = 2 * Ld / 5
+    A = omega * r_in / (2 * (r_in - r_out))
+    B = - omega * r_in * r_out / (r_in - r_out)
+    C = omega * r_in ** 2 * r_out / (r_in - r_out) / 2
+    psi_expr = conditional(r < r_in,
+                           omega * r ** 2 / 2,
+                           conditional(r < r_out,
+                                       A * r ** 2 + B * r + C,
+                                       A * r_out ** 2 + B * r_out + C))
+    psi.interpolate(psi_expr)
+
+    gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)])
+    u.project(gradperp(psi))
+
+    return stepper, tmax, true_field
+
+
+@pytest.mark.parametrize('space', [None, 'equispaced'])
+def test_zero_limiter(tmpdir, space):
+
+    # Setup and run
+    dirname = str(tmpdir)
+
+    stepper, tmax, true_field = setup_zero_limiter(dirname, space)
+
+    stepper.run(t=0, tmax=tmax)
+
+    final_field = stepper.fields('tracer')
+
+    # Check tracer is roughly in the correct place
+    assert norm(true_field - final_field) / norm(true_field) < 0.05, \
+        'Something appears to have gone wrong with transport of tracer using a limiter'
+
+    # Check for no new overshoots
+    assert np.min(final_field.dat.data) >= 0.0, \
+        'Application of limiter has not prevented negative values'
diff --git a/unit-tests/kernel_tests/test_clip_zero_kernel.py b/unit-tests/kernel_tests/test_clip_zero_kernel.py
new file mode 100644
index 000000000..9a664998a
--- /dev/null
+++ b/unit-tests/kernel_tests/test_clip_zero_kernel.py
@@ -0,0 +1,40 @@
+"""
+A test of the ClipZero kernel, which is used to enforce non-negativity.
+"""
+
+from firedrake import UnitSquareMesh, Function, FunctionSpace
+from gusto import kernels
+import numpy as np
+
+
+def test_clip_zero():
+
+    # ------------------------------------------------------------------------ #
+    # Set up meshes and spaces
+    # ------------------------------------------------------------------------ #
+
+    mesh = UnitSquareMesh(3, 3)
+
+    DG1 = FunctionSpace(mesh, "DG", 1)
+
+    field = Function(DG1)
+
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+
+    field.dat.data[:] = -3.0
+
+    # ------------------------------------------------------------------------ #
+    # Apply kernel
+    # ------------------------------------------------------------------------ #
+
+    kernel = kernels.ClipZero(DG1)
+    kernel.apply(field, field)
+
+    # ------------------------------------------------------------------------ #
+    # Check values
+    # ------------------------------------------------------------------------ #
+
+    assert np.all(field.dat.data >= 0.0), \
+        'ClipZero kernel is not enforcing non-negativity'

From bdcc9429d08ae837fdd330413c53e819565f5dd2 Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Tue, 27 Feb 2024 21:05:02 +0000
Subject: [PATCH 32/38] remove unused variable

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

diff --git a/gusto/kernels.py b/gusto/kernels.py
index 7445d7a68..8ca72b363 100644
--- a/gusto/kernels.py
+++ b/gusto/kernels.py
@@ -10,7 +10,7 @@
 """
 
 from firedrake import dx
-from firedrake.parloops import par_loop, READ, WRITE, RW
+from firedrake.parloops import par_loop, READ, WRITE
 
 
 class LimitMidpoints():

From 3c1b2f6e60f1ecc9c025a23650c61fd94abf575c Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Tue, 27 Feb 2024 21:35:47 +0000
Subject: [PATCH 33/38] add limiter to an example

---
 examples/compressible/unsaturated_bubble.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py
index 2a1967a79..074224cf8 100644
--- a/examples/compressible/unsaturated_bubble.py
+++ b/examples/compressible/unsaturated_bubble.py
@@ -67,6 +67,7 @@
 VCG1 = FunctionSpace(mesh, "CG", 1)
 Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element())
 Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
+Vt = domain.spaces("theta")
 
 u_opts = RecoveryOptions(embedding_space=Vu_DG1,
                          recovered_space=Vu_CG1,
@@ -92,10 +93,12 @@
 # Physics schemes
 # NB: can't yet use wrapper or limiter options with physics
 rainfall_method = DGUpwind(eqns, 'rain', outflow=True)
+zero_limiter = MixedFSLimiter(eqns, {'water_vapour': ZeroLimiter(Vt),
+                                     'cloud_water': ZeroLimiter(Vt)})
 physics_schemes = [(Fallout(eqns, 'rain', domain, rainfall_method), SSPRK3(domain)),
                    (Coalescence(eqns), ForwardEuler(domain)),
                    (EvaporationOfRain(eqns), ForwardEuler(domain)),
-                   (SaturationAdjustment(eqns), ForwardEuler(domain))]
+                   (SaturationAdjustment(eqns), ForwardEuler(domain, limiter=zero_limiter))]
 
 # Time stepper
 stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields,
@@ -116,7 +119,6 @@
 
 # spaces
 Vu = domain.spaces("HDiv")
-Vt = domain.spaces("theta")
 Vr = domain.spaces("DG")
 x, z = SpatialCoordinate(mesh)
 quadrature_degree = (4, 4)

From 07003818f7c00bfb7f82652541243a4024b7b3dc Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 28 Feb 2024 03:25:54 +0000
Subject: [PATCH 34/38] set up original spaces for MixedFSWrapper within the
 subwrappers

---
 gusto/time_discretisation.py |  6 +-----
 gusto/wrappers.py            | 27 ++++++++++++---------------
 2 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 5d9297c75..2f6cc009e 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -184,11 +184,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                         raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
 
                     field_idx = equation.field_names.index(field)
-
-                    # Store the original space of the tracer
-                    subwrapper.tracer_fs = self.equation.spaces[field_idx]
-
-                    subwrapper.setup()
+                    subwrapper.setup(self.equation.spaces[field_idx])
 
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 0ac2728c3..670fcacd1 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -33,7 +33,6 @@ def __init__(self, time_discretisation, wrapper_options):
         self.time_discretisation = time_discretisation
         self.options = wrapper_options
         self.solver_parameters = None
-        self.tracer_fs = None
 
     @abstractmethod
     def setup(self):
@@ -77,7 +76,7 @@ class EmbeddedDGWrapper(Wrapper):
     the original space.
     """
 
-    def setup(self):
+    def setup(self, previous_space=None):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, EmbeddedDGOptions), \
@@ -86,8 +85,8 @@ def setup(self):
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if self.tracer_fs is not None:
-            original_space = self.tracer_fs
+        if previous_space is not None:
+            original_space = previous_space
         else:
             original_space = self.time_discretisation.fs
 
@@ -110,8 +109,8 @@ def setup(self):
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
 
-        if self.tracer_fs is not None:
-            self.x_projected = Function(self.tracer_fs)
+        if previous_space is not None:
+            self.x_projected = Function(previous_space)
         elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
         else:
@@ -166,19 +165,17 @@ class RecoveryWrapper(Wrapper):
     field is then returned to the original space.
     """
 
-    def setup(self):
+    def setup(self, previous_space=None):
         """Sets up function spaces and fields needed for this wrapper."""
 
-        print(self.options)
-
         assert isinstance(self.options, RecoveryOptions), \
             'Recovery wrapper can only be used with Recovery Options'
 
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if self.tracer_fs is not None:
-            original_space = self.tracer_fs
+        if previous_space is not None:
+            original_space = previous_space
         else:
             original_space = self.time_discretisation.fs
 
@@ -198,16 +195,16 @@ def setup(self):
         # Internal variables to be used
         # -------------------------------------------------------------------- #
 
-        if self.tracer_fs is not None:
-            self.x_in_tmp = Function(self.tracer_fs)
+        if previous_space is not None:
+            self.x_in_tmp = Function(previous_space)
         else:
             self.x_in_tmp = Function(self.time_discretisation.fs)
 
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
 
-        if self.tracer_fs is not None:
-            self.x_projected = Function(self.tracer_fs)
+        if previous_space is not None:
+            self.x_projected = Function(previous_space)
         elif self.time_discretisation.idx is None:
             self.x_projected = Function(equation.function_space)
         else:

From 2f517e1df8f06b93c6b7f6811728e85d3cc94dac Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 28 Feb 2024 19:30:26 +0000
Subject: [PATCH 35/38] make original_space a base wrapper property

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

diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 670fcacd1..c4867e37b 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -33,14 +33,27 @@ def __init__(self, time_discretisation, wrapper_options):
         self.time_discretisation = time_discretisation
         self.options = wrapper_options
         self.solver_parameters = None
+        self.original_space = None
 
     @abstractmethod
-    def setup(self):
+    def setup(self, original_space=None):
         """
-        Performs standard set up routines, and is to be called by the setup
-        method of the underlying time discretisation.
+        Store the original function space of the prognostic variable.
+
+        Within each child wrapper, setup performs standard set up routines,
+        and is to be called by the setup method of the underlying
+        time discretisation.
+
+        Args:
+            original_space (:class:`FunctionSpace`): the space that the
+        prognostic variable is defined on. This is a subset space of
+        a mixed function space when using a MixedFSWrapper. Defaults
+        to None, when not using a MixedFSWrapper.
         """
-        pass
+        if original_space is not None:
+            self.original_space = original_space
+        else:
+            self.original_space = self.time_discretisation.fs
 
     @abstractmethod
     def pre_apply(self):
@@ -76,26 +89,23 @@ class EmbeddedDGWrapper(Wrapper):
     the original space.
     """
 
-    def setup(self, previous_space=None):
+    def setup(self, original_space=None):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, EmbeddedDGOptions), \
             'Embedded DG wrapper can only be used with Embedded DG Options'
 
+        super().setup(original_space)
+
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if previous_space is not None:
-            original_space = previous_space
-        else:
-            original_space = self.time_discretisation.fs
-
         # -------------------------------------------------------------------- #
         # Set up spaces to be used with wrapper
         # -------------------------------------------------------------------- #
 
         if self.options.embedding_space is None:
-            V_elt = BrokenElement(original_space.ufl_element())
+            V_elt = BrokenElement(self.original_space.ufl_element())
             self.function_space = FunctionSpace(domain.mesh, V_elt)
         else:
             self.function_space = self.options.embedding_space
@@ -109,10 +119,8 @@ def setup(self, previous_space=None):
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
 
-        if previous_space is not None:
-            self.x_projected = Function(previous_space)
-        elif self.time_discretisation.idx is None:
-            self.x_projected = Function(equation.function_space)
+        if self.time_discretisation.idx is None:
+            self.x_projected = Function(self.original_space)
         else:
             self.x_projected = Function(equation.spaces[self.time_discretisation.idx])
 
@@ -165,26 +173,23 @@ class RecoveryWrapper(Wrapper):
     field is then returned to the original space.
     """
 
-    def setup(self, previous_space=None):
+    def setup(self, original_space=None):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, RecoveryOptions), \
             'Recovery wrapper can only be used with Recovery Options'
 
+        super().setup(original_space)
+
         domain = self.time_discretisation.domain
         equation = self.time_discretisation.equation
 
-        if previous_space is not None:
-            original_space = previous_space
-        else:
-            original_space = self.time_discretisation.fs
-
         # -------------------------------------------------------------------- #
         # Set up spaces to be used with wrapper
         # -------------------------------------------------------------------- #
 
         if self.options.embedding_space is None:
-            V_elt = BrokenElement(original_space.ufl_element())
+            V_elt = BrokenElement(self.original_space.ufl_element())
             self.function_space = FunctionSpace(domain.mesh, V_elt)
         else:
             self.function_space = self.options.embedding_space
@@ -195,18 +200,12 @@ def setup(self, previous_space=None):
         # Internal variables to be used
         # -------------------------------------------------------------------- #
 
-        if previous_space is not None:
-            self.x_in_tmp = Function(previous_space)
-        else:
-            self.x_in_tmp = Function(self.time_discretisation.fs)
-
+        self.x_in_tmp = Function(self.original_space)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space)
 
-        if previous_space is not None:
-            self.x_projected = Function(previous_space)
-        elif self.time_discretisation.idx is None:
-            self.x_projected = Function(equation.function_space)
+        if self.time_discretisation.idx is None:
+            self.x_projected = Function(self.original_space)
         else:
             self.x_projected = Function(equation.spaces[self.time_discretisation.idx])
 

From 5c2980dbf9b7339c9e65651a5faa70af1a566892 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Thu, 29 Feb 2024 19:09:37 +0000
Subject: [PATCH 36/38] make original_space a required argument for
 wrapper.setup()

---
 gusto/time_discretisation.py |  7 +++++--
 gusto/wrappers.py            | 14 +++++---------
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 2f6cc009e..9a30e337d 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -184,7 +184,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                         raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
 
                     field_idx = equation.field_names.index(field)
-                    subwrapper.setup(self.equation.spaces[field_idx])
+                    subwrapper.setup(equation.spaces[field_idx])
 
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
@@ -200,7 +200,10 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     map_if_true=replace_test_function(new_test_mixed))
 
             else:
-                self.wrapper.setup()
+                if self.wrapper_name == "supg":
+                    self.wrapper.setup()
+                else:
+                    self.wrapper.setup(self.fs)
                 self.fs = self.wrapper.function_space
                 if self.solver_parameters is None:
                     self.solver_parameters = self.wrapper.solver_parameters
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index c4867e37b..f3fd54c2f 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -36,7 +36,7 @@ def __init__(self, time_discretisation, wrapper_options):
         self.original_space = None
 
     @abstractmethod
-    def setup(self, original_space=None):
+    def setup(self, original_space):
         """
         Store the original function space of the prognostic variable.
 
@@ -47,13 +47,9 @@ def setup(self, original_space=None):
         Args:
             original_space (:class:`FunctionSpace`): the space that the
         prognostic variable is defined on. This is a subset space of
-        a mixed function space when using a MixedFSWrapper. Defaults
-        to None, when not using a MixedFSWrapper.
+        a mixed function space when using a MixedFSWrapper. 
         """
-        if original_space is not None:
-            self.original_space = original_space
-        else:
-            self.original_space = self.time_discretisation.fs
+        self.original_space = original_space
 
     @abstractmethod
     def pre_apply(self):
@@ -89,7 +85,7 @@ class EmbeddedDGWrapper(Wrapper):
     the original space.
     """
 
-    def setup(self, original_space=None):
+    def setup(self, original_space):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, EmbeddedDGOptions), \
@@ -173,7 +169,7 @@ class RecoveryWrapper(Wrapper):
     field is then returned to the original space.
     """
 
-    def setup(self, original_space=None):
+    def setup(self, original_space):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, RecoveryOptions), \

From 222c396edeeac69e5f5465c4a2ca4d3904e45b56 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Thu, 29 Feb 2024 19:11:25 +0000
Subject: [PATCH 37/38] lint

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

diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index f3fd54c2f..3ca9c4347 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -47,7 +47,7 @@ def setup(self, original_space):
         Args:
             original_space (:class:`FunctionSpace`): the space that the
         prognostic variable is defined on. This is a subset space of
-        a mixed function space when using a MixedFSWrapper. 
+        a mixed function space when using a MixedFSWrapper.
         """
         self.original_space = original_space
 

From dd5ab84767c5be5043a258c0a1e042b43f134f2c Mon Sep 17 00:00:00 2001
From: Tom Bendall <thomas.bendall@metoffice.gov.uk>
Date: Mon, 4 Mar 2024 14:30:44 +0000
Subject: [PATCH 38/38] pick up changes to matrix assembly from Firedrake

---
 gusto/preconditioners.py | 39 +++++++++++++++++++--------------------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/gusto/preconditioners.py b/gusto/preconditioners.py
index b400e2898..a471b18bc 100644
--- a/gusto/preconditioners.py
+++ b/gusto/preconditioners.py
@@ -9,6 +9,7 @@
 from gusto.recovery.recovery_kernels import AverageKernel, AverageWeightings
 from pyop2.profiling import timed_region, timed_function
 from pyop2.utils import as_tuple
+from functools import partial
 
 
 __all__ = ["VerticalHybridizationPC"]
@@ -49,8 +50,7 @@ def initialize(self, pc):
                                FiniteElement, TensorProductElement,
                                TrialFunction, TrialFunctions, TestFunction,
                                DirichletBC, interval, MixedElement, BrokenElement)
-        from firedrake.assemble import (allocate_matrix, OneFormAssembler,
-                                        TwoFormAssembler)
+        from firedrake.assemble import get_assembler
         from firedrake.formmanipulation import split_form
         from ufl.algorithms.replace import replace
         from ufl.cell import TensorProductCell
@@ -176,22 +176,21 @@ def initialize(self, pc):
 
         # Assemble the Schur complement operator and right-hand side
         self.schur_rhs = Cofunction(Vv_tr.dual())
-        self._assemble_Srhs = OneFormAssembler(
+        self._assemble_Srhs = partial(get_assembler(
             K * Atilde.inv * AssembledVector(self.broken_residual),
-            tensor=self.schur_rhs,
-            form_compiler_parameters=self.ctx.fc_params).assemble
+            form_compiler_parameters=self.ctx.fc_params).assemble, tensor=self.schur_rhs)
 
         mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
 
         schur_comp = K * Atilde.inv * K.T
-        self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
-                                 form_compiler_parameters=self.ctx.fc_params,
-                                 mat_type=mat_type,
-                                 options_prefix=prefix)
-        self._assemble_S = TwoFormAssembler(schur_comp,
-                                            tensor=self.S,
-                                            bcs=trace_bcs,
-                                            form_compiler_parameters=self.ctx.fc_params).assemble
+        self.S = get_assembler(schur_comp, bcs=trace_bcs,
+                               form_compiler_parameters=self.ctx.fc_params,
+                               mat_type=mat_type,
+                               options_prefix=prefix).allocate()
+        self._assemble_S = partial(get_assembler(
+            schur_comp,
+            bcs=trace_bcs,
+            form_compiler_parameters=self.ctx.fc_params).assemble, tensor=self.S)
 
         self._assemble_S()
         Smat = self.S.petscmat
@@ -228,7 +227,7 @@ def _reconstruction_calls(self, split_mixed_op, split_trace_op):
             split_trace_op (dict): a ``dict`` of split forms that make up the
                 trace contribution in the hybridized mixed system.
         """
-        from firedrake.assemble import OneFormAssembler
+        from firedrake.assemble import get_assembler
 
         # We always eliminate the velocity block first
         id0, id1 = (self.vidx, self.pidx)
@@ -256,15 +255,15 @@ def _reconstruction_calls(self, split_mixed_op, split_trace_op):
         R = K_1.T - C * A.inv * K_0.T
         u_rec = M.solve(f - C * A.inv * g - R * lambdar,
                         decomposition="PartialPivLU")
-        self._sub_unknown = OneFormAssembler(u_rec,
-                                             tensor=u,
-                                             form_compiler_parameters=self.ctx.fc_params).assemble
+        self._sub_unknown = partial(get_assembler(
+            u_rec,
+            form_compiler_parameters=self.ctx.fc_params).assemble, tensor=u)
 
         sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar,
                             decomposition="PartialPivLU")
-        self._elim_unknown = OneFormAssembler(sigma_rec,
-                                              tensor=sigma,
-                                              form_compiler_parameters=self.ctx.fc_params).assemble
+        self._elim_unknown = partial(get_assembler(
+            sigma_rec,
+            form_compiler_parameters=self.ctx.fc_params).assemble, tensor=sigma)
 
     @timed_function("VertHybridRecon")
     def _reconstruct(self):