From ae2afe208c8a832668ef80ea1071555dd99c25b9 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Sat, 28 Dec 2024 20:49:22 +0000
Subject: [PATCH 1/7] start mean mixing ratio augmentation

---
 gusto/spatial_methods/augmentation.py         |  24 +++
 .../model/test_mean_mixing_ratio.py           | 175 ++++++++++++++++++
 2 files changed, 199 insertions(+)
 create mode 100644 integration-tests/model/test_mean_mixing_ratio.py

diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index 42aedaa64..8e5fbc52e 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -236,3 +236,27 @@ def update(self, x_in_mixed):
         logger.debug('Vorticity solve')
         self.solver.solve()
         self.x_in.subfunctions[1].assign(self.Z_in)
+
+
+#class MeanMixingRatio(Augmentation):
+    """
+    This augments a transport problem involving a mixing ratio to
+    include a mean mixing ratio field. This enables posivity to be 
+    ensured during conservative transport.
+
+    Args:
+        domain (:class:`Domain`): The domain object.
+        eqns (:class:`PrognosticEquationSet`): The overarching equation set.
+        mixing_ratio (:class: list): List of mixing ratios that 
+        are to have augmented mean mixing ratio fields.
+        OR, keep as a single mixing ratio, but define 
+        multiple augmentations?
+    """
+
+ #   def __init__(
+ #           self, domain, eqns, mX_name
+ #   ):
+
+
+        # Extract the mixing ratio in question:
+ #       field_idx = self.eqns.field_names.index(mX_name)
diff --git a/integration-tests/model/test_mean_mixing_ratio.py b/integration-tests/model/test_mean_mixing_ratio.py
new file mode 100644
index 000000000..35bdf2c75
--- /dev/null
+++ b/integration-tests/model/test_mean_mixing_ratio.py
@@ -0,0 +1,175 @@
+"""
+Tests the mean mixing ratio augmentation. Here a mixing ratio is transported
+conservatively along with the dry density. The mean mixing ratio augmentation
+should ensure that the mixing ratio remains non-negative. To test this,
+a physics scheme of a sink is applied, which without the mean mixing ratio
+will lead to negative values. Hence, a check is performed that the 
+mixing ratio has no negative values after a couple of time steps.
+"""
+
+from gusto import *
+from firedrake import (
+    PeriodicIntervalMesh, ExtrudedMesh, exp, cos, sin, SpatialCoordinate,
+    assemble, dx, FunctionSpace, pi, min_value, as_vector, BrokenElement,
+    errornorm
+)
+import pytest
+
+
+def setup_mean_mixing_ratio(dirname, pair_of_spaces):
+
+    # Domain
+    Lx = 2000.
+    Hz = 2000.
+
+    # Time parameters
+    dt = 2.
+    tmax = 2000.
+
+    nlayers = 10.  # horizontal layers
+    columns = 10.  # number of columns
+
+    # Define the spaces for the tracers
+    if pair_of_spaces == 'same':
+        rho_d_space = 'DG'
+        m_X_space = 'DG'
+    else:
+        rho_d_space = 'DG'
+        m_X_space = 'theta'
+
+    period_mesh = PeriodicIntervalMesh(columns, Lx)
+    mesh = ExtrudedMesh(period_mesh, layers=nlayers, layer_height=Hz/nlayers)
+    domain = Domain(mesh, dt, "CG", 1)
+    x, z = SpatialCoordinate(mesh)
+
+    V_rho = domain.spaces(rho_d_space)
+    V_m_X = domain.spaces(m_X_space)
+
+    m_X = ActiveTracer(name='m_X', space=m_X_space,
+                       variable_type=TracerVariableType.mixing_ratio,
+                       transport_eqn=TransportEquationType.tracer_conservative,
+                       density_name='rho_d')
+
+    rho_d = ActiveTracer(name='rho_d', space=rho_d_space,
+                         variable_type=TracerVariableType.density,
+                         transport_eqn=TransportEquationType.conservative)
+
+    # Define m_X first to test that the tracers will be
+    # automatically re-ordered such that the density field
+    # is indexed before the mixing ratio.
+    tracers = [m_X, rho_d]
+
+    # Equation
+    V = domain.spaces("HDiv")
+    eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V)
+
+    # IO
+    output = OutputParameters(dirname=dirname)
+    io = IO(domain, output)
+
+    
+    if pair_of_spaces == 'diff':
+        Vt_brok = FunctionSpace(mesh, BrokenElement(V_m_X.ufl_element()))
+        suboptions = {
+            'rho_d': EmbeddedDGOptions(embedding_space=Vt_brok),
+            'm_X': ConservativeEmbeddedDGOptions(
+                rho_name='rho_d',
+                orig_rho_space=V_rho
+            )
+        }
+    else:
+        suboptions = {}
+
+    opts = MixedFSOptions(suboptions=suboptions)
+
+    transport_scheme = SSPRK3(
+        domain, options=opts, rk_formulation=RungeKuttaFormulation.predictor
+    )
+    transport_methods = [DGUpwind(eqn, "m_X"), DGUpwind(eqn, "rho_d")]
+
+    #physics_schemes = [(SourceSink(eqn, 'm_X', -Constant(0.1)), SSPRK3(domain))]
+
+    # Timestepper
+    time_varying = True
+    #stepper = SplitPrescribedTransport(
+    #    eqn, transport_scheme,
+    #    io, time_varying, transport_methods,
+    #    physics_schemes=physics_schemes
+    #)
+    stepper = PrescribedTransport(
+        eqn, transport_scheme, io, time_varying, transport_methods
+    )
+
+    # Initial Conditions
+    # Specify locations of the two Gaussians
+    xc1 = 5.*Lx/8.
+    zc1 = Hz/2.
+
+    xc2 = 3.*Lx/8.
+    zc2 = Hz/2.
+
+    def l2_dist(xc, zc):
+        return min_value(abs(x-xc), Lx-abs(x-xc))**2 + (z-zc)**2
+
+    lc = 2.*Lx/25.
+    m0 = 0.02
+
+    # Set the initial state with two Gaussians for the density
+    # and a linear variation in mixing ratio.
+
+    f0 = 0.5
+    rho_b = 0.5
+
+    g1 = f0*exp(-l2_dist(xc1, zc1)/(lc**2))
+    g2 = f0*exp(-l2_dist(xc2, zc2)/(lc**2))
+
+    rho_d_0 = rho_b + g1 + g2
+
+    # Constant mass field, starting with no mixing
+    # ratio at z=0 and m=0.5 at the model top
+    m_X_0 = m0 + 0.5*z/Hz
+
+    # Set up the divergent, time-varying, velocity field
+    U = Lx/tmax
+    W = U/10.
+
+    def u_t(t):
+        xd = x - U*t
+        u = U - (W*pi*Lx/Hz)*cos(pi*t/tmax)*cos(2*pi*xd/Lx)*cos(pi*z/Hz)
+        w = 2*pi*W*cos(pi*t/tmax)*sin(2*pi*xd/Lx)*sin(pi*z/Hz)
+
+        u_expr = as_vector((u, w))
+
+        return u_expr
+
+    stepper.setup_prescribed_expr(u_t)
+
+    stepper.fields("m_X").interpolate(m_X_0)
+    stepper.fields("rho_d").interpolate(rho_d_0)
+    stepper.fields("u").project(u_t(0))
+
+    m_X_init = Function(V_m_X)
+    rho_d_init = Function(V_rho)
+
+    m_X_init.assign(stepper.fields("m_X"))
+    rho_d_init.assign(stepper.fields("rho_d"))
+
+    return stepper, m_X_init, rho_d_init
+
+
+@pytest.mark.parametrize("pair_of_spaces", ["same", "diff"])
+def test_mean_mixing_ratio(tmpdir, pair_of_spaces):
+
+    # Setup and run
+    dirname = str(tmpdir)
+
+    stepper, m_X_0, rho_d_0 = \
+        setup_mean_mixing_ratio(dirname, pair_of_spaces)
+
+    # Run for two timesteps
+    stepper.run(t=0, tmax=4)
+    m_X = stepper.fields("m_X")
+    rho_d = stepper.fields("rho_d")
+
+    # Check that the mixing ratio has no negative values
+    assert all(m_X_0 >= 0.0), "mean mixing ratio field has not ensured non-negativity"
\ No newline at end of file

From fa1062dc6163b2e3a313252e79058ba1fbb54cd1 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Sun, 29 Dec 2024 16:29:19 +0000
Subject: [PATCH 2/7] create new mixed space for mean mixing ratio

---
 gusto/spatial_methods/augmentation.py | 124 ++++++++++++++++++++++++--
 1 file changed, 117 insertions(+), 7 deletions(-)

diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index 8e5fbc52e..5a5807d7e 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -11,10 +11,12 @@
     transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div,
     Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC
 )
-from firedrake.fml import subject
+from firedrake.fml import (
+    subject, all_terms, replace_subject, keep, replace_test_function
+)
 from gusto import (
     time_derivative, transport, transporting_velocity, TransportEquationType,
-    logger
+    logger, prognostic, mass_weighted
 )
 
 
@@ -238,7 +240,7 @@ def update(self, x_in_mixed):
         self.x_in.subfunctions[1].assign(self.Z_in)
 
 
-#class MeanMixingRatio(Augmentation):
+class MeanMixingRatio(Augmentation):
     """
     This augments a transport problem involving a mixing ratio to
     include a mean mixing ratio field. This enables posivity to be 
@@ -251,12 +253,120 @@ def update(self, x_in_mixed):
         are to have augmented mean mixing ratio fields.
         OR, keep as a single mixing ratio, but define 
         multiple augmentations?
+        START with just a single.
     """
 
- #   def __init__(
- #           self, domain, eqns, mX_name
- #   ):
+    def __init__(
+            self, domain, eqns, mX_name
+    ):
+
+        self.eqns = eqns
+        exist_spaces = eqns.spaces
 
+        self.idx_orig = len(exist_spaces)
+        mean_idx = self.idx_orig
 
         # Extract the mixing ratio in question:
- #       field_idx = self.eqns.field_names.index(mX_name)
+        mX_idx = eqns.field_names.index(mX_name)
+
+        # Define the mean mixing ratio on the DG0 space
+        DG0 = FunctionSpace(domain.mesh, "DG", 0)
+
+         # Set up the scheme for the mean mixing ratio
+
+        mean_mX = Function(DG0)
+        mean_space = DG0
+        exist_spaces.append(mean_space)
+
+        self.fs = MixedFunctionSpace(exist_spaces)
+        self.X = Function(self.fs)
+        self.tests = TestFunctions(self.fs)
+
+        self.x_in = Function(self.fs)
+        self.x_out = Function(self.fs)
+
+        # Compute the new mean mass weighted term,
+        # IF this is conservatively transported.
+        mX_idx = eqns.field_names.index(mX_name)
+
+        #mean_mass = 
+
+        # Now, extract the existing residual:
+        old_residual = eqns.residual
+
+        # Extract terms relating to the mixing ratio of interest
+        mX_residual = old_residual.label_map(
+            lambda t: t.get(prognostic) == mX_name,
+            map_if_true=keep
+        )
+
+        # Replace trial and test functions with the new mixed
+        # function space
+        # Does this work is the subject is mass-weighted???
+        for idx in range(self.idx_orig):
+            field = eqns.field_names[idx]
+            # Seperate logic if mass-weighted or not?
+            print(idx)
+            print(field)
+
+
+            old_residual = old_residual.label_map(
+                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+                map_if_true=replace_subject(self.X[idx], old_idx=idx)
+            )
+            old_residual = old_residual.label_map(
+                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+                map_if_true=replace_test_function(self.tests[idx], old_idx=idx)
+            )
+    
+        # Define the mean mixing ratio residual
+        mean_residual = mX_residual.label_map(
+            lambda t: t.has_label(mass_weighted),
+            #map_if_true=replace_subject(mean_mass, old_idx=mX_idx),
+            map_if_false=replace_subject(self.X[mean_idx], old_idx = mX_idx)
+        )
+
+        mean_residual = mean_residual.label_map(
+            lambda t: t.has_label(mass_weighted),
+            map_if_false=replace_test_function(self.tests[mean_idx], old_idx=mX_idx)
+        )
+
+        self.residual = old_residual + mean_residual
+
+        self.bcs = []
+
+
+    def pre_apply(self, x_in):
+        """
+        Sets the original fields, i.e. not the mean mixing ratios
+
+        Args:
+            x_in (:class:`Function`): The input fields
+        """
+
+        #for idx, field in enumerate(self.eqn.field_names):
+        #    self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx])
+        for idx in range(self.idx_orig):
+            self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx])
+
+
+    def post_apply(self, x_out):
+        """
+        Sets the output fields, i.e. not the mean mixing ratios
+
+        Args:
+            x_out (:class:`Function`): The output fields
+        """
+
+        for idx, field in enumerate(self.eqn.field_names):
+            x_out.subfunctions[idx].assign(self.x_out.subfunctions[idx])
+
+    def update(self, x_in_mixed):
+        """
+        ...
+
+        Args:
+            x_in_mixed (:class:`Function`): The mixed function to update.
+        """
+        
+        pass
\ No newline at end of file

From e9b05a285f58ea1d8885083518a41100d384c822 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Thu, 2 Jan 2025 16:25:44 +0000
Subject: [PATCH 3/7] set the mean mixing ratio augmentation residual

---
 gusto/spatial_methods/augmentation.py | 107 ++++++++++++++++++++++----
 1 file changed, 93 insertions(+), 14 deletions(-)

diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index 5a5807d7e..9734bf21a 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -12,7 +12,8 @@
     Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC
 )
 from firedrake.fml import (
-    subject, all_terms, replace_subject, keep, replace_test_function
+    subject, all_terms, replace_subject, keep, replace_test_function,
+    replace_trial_function, drop
 )
 from gusto import (
     time_derivative, transport, transporting_velocity, TransportEquationType,
@@ -274,7 +275,7 @@ def __init__(
 
          # Set up the scheme for the mean mixing ratio
 
-        mean_mX = Function(DG0)
+        mean_mX = Function(DG0, name='mean_mX')
         mean_space = DG0
         exist_spaces.append(mean_space)
 
@@ -282,6 +283,10 @@ def __init__(
         self.X = Function(self.fs)
         self.tests = TestFunctions(self.fs)
 
+        print(self.X)
+
+        self.bcs = []
+
         self.x_in = Function(self.fs)
         self.x_out = Function(self.fs)
 
@@ -289,17 +294,91 @@ def __init__(
         # IF this is conservatively transported.
         mX_idx = eqns.field_names.index(mX_name)
 
-        #mean_mass = 
+        old_residual = eqns.residual
+
+        mean_residual = old_residual.label_map(
+            lambda t: t.get(prognostic) == mX_name,
+            map_if_false=drop
+        )
+
+        # Replace trial functions with those in the new mixed function space
+        for term in eqns.residual:
+            print('\n')
+            print(term.form)
 
+        for idx in range(self.idx_orig):
+            field = eqns.field_names[idx]
+            # Seperate logic if mass-weighted or not?
+            print(idx)
+            print(field)
+
+            prog = split(self.X)[idx]
+
+            print('\n residual before change')
+            print(old_residual.form)
+            old_residual = old_residual.label_map(
+                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+                map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
+            )
+            old_residual = old_residual.label_map(
+                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+                map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
+            )
+            print('\n residual after change')
+            print(old_residual.form)
+
+            #old_residual = old_residual.label_map(
+            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+            #    map_if_true=replace_subject(self.X[idx], old_idx=idx)
+            #)
+            #old_residual = old_residual.label_map(
+            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+            #    map_if_true=replace_test_function(self.tests[idx], old_idx=idx)
+            #)
+    
+        # Define the mean mixing ratio residual
+        #mean_residual = mX_residual.label_map(
+        #    lambda t: t.has_label(mass_weighted),
+            #map_if_true=replace_subject(mean_mass, old_idx=mX_idx),
+        #    map_if_false=replace_subject(self.X[mean_idx], old_idx = mX_idx)
+        #)
+
+        print('\n mean mX residual before change')
+        print(mean_residual.form)
+        mean_residual = mean_residual.label_map(
+            all_terms,
+            replace_subject(self.X, old_idx=mX_idx, new_idx = mean_idx)
+        )
+        mean_residual = mean_residual.label_map(
+            all_terms,
+            replace_test_function(self.tests, old_idx=mX_idx, new_idx = mean_idx)
+        )
+        print('\n mean mX residual after change')
+        print(mean_residual.form)
+
+        # Form the new residual
+        self.residual = old_residual + mean_residual
+
+        print('\n full residual')
+        print(self.residual.form)
+
+        
+
+    def setup_residual(self):
+        # Update the residual
         # Now, extract the existing residual:
         old_residual = eqns.residual
 
+        print(old_residual.form)
+
         # Extract terms relating to the mixing ratio of interest
         mX_residual = old_residual.label_map(
             lambda t: t.get(prognostic) == mX_name,
-            map_if_true=keep
+            map_if_false=drop
         )
 
+        print(mX_residual.form)
+
         # Replace trial and test functions with the new mixed
         # function space
         # Does this work is the subject is mass-weighted???
@@ -309,15 +388,16 @@ def __init__(
             print(idx)
             print(field)
 
+            prog = split(self.X)[idx]
 
-            old_residual = old_residual.label_map(
-                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-                map_if_true=replace_subject(self.X[idx], old_idx=idx)
-            )
-            old_residual = old_residual.label_map(
-                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-                map_if_true=replace_test_function(self.tests[idx], old_idx=idx)
-            )
+            #old_residual = old_residual.label_map(
+            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+            #    map_if_true=replace_subject(self.X[idx], old_idx=idx)
+            #)
+            #old_residual = old_residual.label_map(
+            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+            #    map_if_true=replace_test_function(self.tests[idx], old_idx=idx)
+            #)
     
         # Define the mean mixing ratio residual
         mean_residual = mX_residual.label_map(
@@ -331,10 +411,9 @@ def __init__(
             map_if_false=replace_test_function(self.tests[mean_idx], old_idx=mX_idx)
         )
 
+        # Form the new residual
         self.residual = old_residual + mean_residual
 
-        self.bcs = []
-
 
     def pre_apply(self, x_in):
         """

From 9adf64e223779fa357e435a184bfe6d339959706 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Fri, 3 Jan 2025 15:42:27 +0000
Subject: [PATCH 4/7] more augmentation changes

---
 gusto/spatial_methods/augmentation.py   | 77 +++++++++----------------
 gusto/timestepping/split_timestepper.py | 10 ++++
 2 files changed, 36 insertions(+), 51 deletions(-)

diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index 9734bf21a..c0d59098b 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -19,6 +19,7 @@
     time_derivative, transport, transporting_velocity, TransportEquationType,
     logger, prognostic, mass_weighted
 )
+import copy
 
 
 class Augmentation(object, metaclass=ABCMeta):
@@ -76,6 +77,8 @@ def __init__(
             self, domain, eqns, transpose_commutator=True, supg=False
     ):
 
+        self.name = 'vorticity'
+
         V_vel = domain.spaces('HDiv')
         V_vort = domain.spaces('H1')
 
@@ -261,6 +264,11 @@ def __init__(
             self, domain, eqns, mX_name
     ):
 
+        self.name = 'mean_mixing_ratio'
+        self.mX_name = mX_name
+        self.mean_name = 'mean_'+mX_name
+        print(self.mean_name)
+
         self.eqns = eqns
         exist_spaces = eqns.spaces
 
@@ -275,7 +283,7 @@ def __init__(
 
          # Set up the scheme for the mean mixing ratio
 
-        mean_mX = Function(DG0, name='mean_mX')
+        mean_mX = Function(DG0, name=self.mean_name)
         mean_space = DG0
         exist_spaces.append(mean_space)
 
@@ -300,6 +308,7 @@ def __init__(
             lambda t: t.get(prognostic) == mX_name,
             map_if_false=drop
         )
+        mean_residual = prognostic.update_value(mean_residual, self.mean_name)
 
         # Replace trial functions with those in the new mixed function space
         for term in eqns.residual:
@@ -362,58 +371,24 @@ def __init__(
         print('\n full residual')
         print(self.residual.form)
 
-        
-
-    def setup_residual(self):
-        # Update the residual
-        # Now, extract the existing residual:
-        old_residual = eqns.residual
-
-        print(old_residual.form)
-
-        # Extract terms relating to the mixing ratio of interest
-        mX_residual = old_residual.label_map(
-            lambda t: t.get(prognostic) == mX_name,
-            map_if_false=drop
-        )
-
-        print(mX_residual.form)
-
-        # Replace trial and test functions with the new mixed
-        # function space
-        # Does this work is the subject is mass-weighted???
-        for idx in range(self.idx_orig):
-            field = eqns.field_names[idx]
-            # Seperate logic if mass-weighted or not?
-            print(idx)
-            print(field)
-
-            prog = split(self.X)[idx]
-
-            #old_residual = old_residual.label_map(
-            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-            #    map_if_true=replace_subject(self.X[idx], old_idx=idx)
-            #)
-            #old_residual = old_residual.label_map(
-            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-            #    map_if_true=replace_test_function(self.tests[idx], old_idx=idx)
-            #)
-    
-        # Define the mean mixing ratio residual
-        mean_residual = mX_residual.label_map(
-            lambda t: t.has_label(mass_weighted),
-            #map_if_true=replace_subject(mean_mass, old_idx=mX_idx),
-            map_if_false=replace_subject(self.X[mean_idx], old_idx = mX_idx)
-        )
+        print('yoyoyoy')
 
-        mean_residual = mean_residual.label_map(
-            lambda t: t.has_label(mass_weighted),
-            map_if_false=replace_test_function(self.tests[mean_idx], old_idx=mX_idx)
-        )
-
-        # Form the new residual
-        self.residual = old_residual + mean_residual
+        for term in self.residual:
+            print(term.get(prognostic))
 
+    def setup_transport(self, spatial_methods):
+        # Copy spatial method for the mixing ratio onto the 
+        # mean mixing ratio.
+        mX_spatial_method = next(method for method in spatial_methods if method.variable == self.mX_name)
+        
+        mean_spatial_method = copy.copy(mX_spatial_method)
+        mean_spatial_method.variable = self.mean_name
+        self.spatial_methods = copy.copy(spatial_methods)
+        self.spatial_methods.append(mean_spatial_method)
+        for method in self.spatial_methods:
+            print(method.variable)
+            #method.equation.residual = self.residual
+            print(len(method.equation.residual))
 
     def pre_apply(self, x_in):
         """
diff --git a/gusto/timestepping/split_timestepper.py b/gusto/timestepping/split_timestepper.py
index 2006c73ac..5d9886436 100644
--- a/gusto/timestepping/split_timestepper.py
+++ b/gusto/timestepping/split_timestepper.py
@@ -296,7 +296,17 @@ def transporting_velocity(self):
         return self.fields('u')
 
     def setup_scheme(self):
+        print('Setting up base equation')
         self.setup_equation(self.equation)
+
+        # If there is an augmentation, set up these transport terms
+        if self.scheme.augmentation is not None:
+            if self.scheme.augmentation.name == 'mean_mixing_ratio':
+                self.scheme.augmentation.setup_transport(self.spatial_methods)
+                print('Setting up augmented equation')
+                self.setup_equation(self.scheme.augmentation)
+
+
         # Go through and label all non-physics terms with a "dynamics" label
         dynamics = Label('dynamics')
         self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics)

From 97b8917f250ee91da31532f98cc8a1678cd5f82d Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Tue, 7 Jan 2025 20:27:03 +0000
Subject: [PATCH 5/7] implementation for a single mixing ratio

---
 gusto/spatial_methods/augmentation.py   | 204 ++++++++++++++++++------
 gusto/timestepping/split_timestepper.py |   4 +-
 2 files changed, 160 insertions(+), 48 deletions(-)

diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index c0d59098b..740afe31d 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -9,7 +9,8 @@
     LinearVariationalProblem, LinearVariationalSolver, lhs, rhs, dot,
     ds_b, ds_v, ds_t, ds, FacetNormal, TestFunction, TrialFunction,
     transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div,
-    Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC
+    Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC,
+    Projector
 )
 from firedrake.fml import (
     subject, all_terms, replace_subject, keep, replace_test_function,
@@ -19,6 +20,7 @@
     time_derivative, transport, transporting_velocity, TransportEquationType,
     logger, prognostic, mass_weighted
 )
+from gusto.spatial_methods.transport_methods import DGUpwind
 import copy
 
 
@@ -269,17 +271,26 @@ def __init__(
         self.mean_name = 'mean_'+mX_name
         print(self.mean_name)
 
-        self.eqns = eqns
+        self.eqn_orig = eqns
+        self.domain=domain
         exist_spaces = eqns.spaces
 
         self.idx_orig = len(exist_spaces)
         mean_idx = self.idx_orig
+        self.mean_idx = mean_idx
 
         # Extract the mixing ratio in question:
         mX_idx = eqns.field_names.index(mX_name)
+        self.mX_idx = mX_idx
 
         # Define the mean mixing ratio on the DG0 space
         DG0 = FunctionSpace(domain.mesh, "DG", 0)
+        self.DG0 = Function(DG0)
+        mX_space = eqns.spaces[mean_idx]
+
+        self.mX_in = Function(mX_space)
+        self.mean_in = Function(mean_space)
+        self.compute_mean = Projector(mX_space, DG0)
 
          # Set up the scheme for the mean mixing ratio
 
@@ -293,48 +304,57 @@ def __init__(
 
         print(self.X)
 
-        self.bcs = []
+        self.bcs = None
 
         self.x_in = Function(self.fs)
         self.x_out = Function(self.fs)
 
         # Compute the new mean mass weighted term,
         # IF this is conservatively transported.
-        mX_idx = eqns.field_names.index(mX_name)
 
-        old_residual = eqns.residual
+        #old_residual = eqns.residual
 
-        mean_residual = old_residual.label_map(
-            lambda t: t.get(prognostic) == mX_name,
-            map_if_false=drop
-        )
-        mean_residual = prognostic.update_value(mean_residual, self.mean_name)
+        #mean_residual = old_residual.label_map(
+        #    lambda t: t.get(prognostic) == mX_name,
+        #    map_if_false=drop
+        #)
+        #mean_residual = prognostic.update_value(mean_residual, self.mean_name)
 
         # Replace trial functions with those in the new mixed function space
-        for term in eqns.residual:
-            print('\n')
-            print(term.form)
+        #for term in eqns.residual:
+        #    print('\n')
+        #    print(term.form)
 
-        for idx in range(self.idx_orig):
-            field = eqns.field_names[idx]
-            # Seperate logic if mass-weighted or not?
-            print(idx)
-            print(field)
-
-            prog = split(self.X)[idx]
+        # If I do this later on with the transport terms, maybe I don't need to do this here?
 
-            print('\n residual before change')
-            print(old_residual.form)
-            old_residual = old_residual.label_map(
-                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-                map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
-            )
-            old_residual = old_residual.label_map(
-                lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-                map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
-            )
-            print('\n residual after change')
-            print(old_residual.form)
+        #for idx in range(self.idx_orig):
+        #    field = eqns.field_names[idx]
+            # Seperate logic if mass-weighted or not?
+        #    print(idx)
+        #    print(field)
+
+        #    prog = split(self.X)[idx]
+
+        #    print('\n residual term before change')
+        #    print(old_residual.label_map(
+        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+        #        map_if_false=drop
+        #    ).form)
+
+        #    old_residual = old_residual.label_map(
+        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+        #        map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
+        #    )
+        #    old_residual = old_residual.label_map(
+        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+        #        map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
+        #    )
+            
+        #    print('\n residual term after change')
+         #   print(old_residual.label_map(
+        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
+        #        map_if_false=drop
+        #    ).form)
 
             #old_residual = old_residual.label_map(
             #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
@@ -352,15 +372,87 @@ def __init__(
         #    map_if_false=replace_subject(self.X[mean_idx], old_idx = mX_idx)
         #)
 
-        print('\n mean mX residual before change')
+        #print('\n mean mX residual before change')
+        #print(mean_residual.form)
+        #mean_residual = mean_residual.label_map(
+        #    all_terms,
+        #    replace_subject(self.X, old_idx=mX_idx, new_idx = mean_idx)
+        #)
+        #mean_residual = mean_residual.label_map(
+        #    all_terms,
+        #    replace_test_function(self.tests, old_idx=mX_idx, new_idx = mean_idx)
+        #)
+        #print('\n mean mX residual after change')
+        #print(mean_residual.form)
+
+        # Form the new residual
+        #self.residual = old_residual + mean_residual
+
+        #print('\n full residual')
+        #print(self.residual.form)
+
+        #print('yoyoyoy')
+
+        #for term in self.residual:
+        #    print(term.get(prognostic))
+
+    def setup_transport(self, spatial_methods, equation):
+        # Copy spatial method for the mixing ratio onto the 
+        # mean mixing ratio.
+
+        old_residual = equation.residual
+
+        # Copy the mean mixing ratio residual terms:
+        mean_residual = old_residual.label_map(
+            lambda t: t.get(prognostic) == self.mX_name,
+            map_if_false=drop
+        )
+        mean_residual = prognostic.update_value(mean_residual, self.mean_name)
+
+        print('\n in setup_transport')
+
+        # Replace the tests and trial functions for all terms
+        # of the fields in the original equation
+        for idx in range(self.idx_orig):
+            field = self.eqn_orig.field_names[idx]
+            # Seperate logic if mass-weighted or not?
+            print('\n', idx)
+            print(field)
+
+            prog = split(self.X)[idx]
+
+            print('\n residual term before change')
+            print(old_residual.label_map(
+                lambda t: t.get(prognostic) == field,
+                map_if_false=drop
+            ).form)
+
+            old_residual = old_residual.label_map(
+                lambda t: t.get(prognostic) == field,
+                map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
+            )
+            old_residual = old_residual.label_map(
+                lambda t: t.get(prognostic) == field,
+                map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
+            )
+            print('\n residual term after change')
+            print(old_residual.label_map(
+                lambda t: t.get(prognostic) == field,
+                map_if_false=drop
+            ).form)
+
+        print('\n now setting up mean mixing ratio residual terms')
+
+
+        print('\n mean mX residual after change')
         print(mean_residual.form)
         mean_residual = mean_residual.label_map(
             all_terms,
-            replace_subject(self.X, old_idx=mX_idx, new_idx = mean_idx)
+            replace_subject(self.X, old_idx=self.mX_idx, new_idx=self.mean_idx)
         )
         mean_residual = mean_residual.label_map(
             all_terms,
-            replace_test_function(self.tests, old_idx=mX_idx, new_idx = mean_idx)
+            replace_test_function(self.tests, old_idx=self.mX_idx, new_idx=self.mean_idx)
         )
         print('\n mean mX residual after change')
         print(mean_residual.form)
@@ -368,17 +460,13 @@ def __init__(
         # Form the new residual
         self.residual = old_residual + mean_residual
 
-        print('\n full residual')
-        print(self.residual.form)
+        #Check these two forms
+        print('\n Original equation with residual of length, ', len(equation.residual))
+        print('\n Augmented equation with residual of length, ', len(self.residual))
 
-        print('yoyoyoy')
 
-        for term in self.residual:
-            print(term.get(prognostic))
 
-    def setup_transport(self, spatial_methods):
-        # Copy spatial method for the mixing ratio onto the 
-        # mean mixing ratio.
+    def setup_transport_old(self, spatial_methods):
         mX_spatial_method = next(method for method in spatial_methods if method.variable == self.mX_name)
         
         mean_spatial_method = copy.copy(mX_spatial_method)
@@ -387,9 +475,19 @@ def setup_transport(self, spatial_methods):
         self.spatial_methods.append(mean_spatial_method)
         for method in self.spatial_methods:
             print(method.variable)
-            #method.equation.residual = self.residual
+            method.equation.residual = self.residual
+            print(method.form.form)
             print(len(method.equation.residual))
 
+        # Alternatively, redo all the spatial methods
+        # using the new mixed function space.
+        # So, want to make a new list of spatial methods
+        new_spatial_methods = []
+        for method in self.spatial_methods:
+            # Determine the tye of transport method:
+            new_method = DGUpwind(self, method.variable)
+            new_spatial_methods.append(new_method)
+
     def pre_apply(self, x_in):
         """
         Sets the original fields, i.e. not the mean mixing ratios
@@ -403,6 +501,12 @@ def pre_apply(self, x_in):
         for idx in range(self.idx_orig):
             self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx])
 
+        # Set the mean mixing ratio to be zero, just because
+        #DG0 = FunctionSpace(self.domain.mesh, "DG", 0)
+        #mean_mX = Function(DG0, name=self.mean_name)
+        
+        #self.x_in.subfunctions[self.mean_idx].assign(mean_mX)
+
 
     def post_apply(self, x_out):
         """
@@ -412,7 +516,7 @@ def post_apply(self, x_out):
             x_out (:class:`Function`): The output fields
         """
 
-        for idx, field in enumerate(self.eqn.field_names):
+        for idx in range(self.idx_orig):
             x_out.subfunctions[idx].assign(self.x_out.subfunctions[idx])
 
     def update(self, x_in_mixed):
@@ -422,5 +526,13 @@ def update(self, x_in_mixed):
         Args:
             x_in_mixed (:class:`Function`): The mixed function to update.
         """
-        
+
+        # Compute the mean mixing ratio
+        # How do I do this?
+        self.mX_in.assign(x_in_mixed[self.mX_idx])
+
+        # Project this into the lowest order space:
+        self.compute_mean.project()
+
+        self.x_in.subfunctions[self.mean_idx].assign(self.mean_out)
         pass
\ No newline at end of file
diff --git a/gusto/timestepping/split_timestepper.py b/gusto/timestepping/split_timestepper.py
index 5d9886436..98b9cb841 100644
--- a/gusto/timestepping/split_timestepper.py
+++ b/gusto/timestepping/split_timestepper.py
@@ -300,11 +300,11 @@ def setup_scheme(self):
         self.setup_equation(self.equation)
 
         # If there is an augmentation, set up these transport terms
+        # Or, perhaps set up the whole residual now ... ?
         if self.scheme.augmentation is not None:
             if self.scheme.augmentation.name == 'mean_mixing_ratio':
-                self.scheme.augmentation.setup_transport(self.spatial_methods)
+                self.scheme.augmentation.setup_transport(self.spatial_methods, self.equation)
                 print('Setting up augmented equation')
-                self.setup_equation(self.scheme.augmentation)
 
 
         # Go through and label all non-physics terms with a "dynamics" label

From bb5cbe133d812994f52a6ff45375095802c67071 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Thu, 9 Jan 2025 17:02:20 +0000
Subject: [PATCH 6/7] first version of mean mixing ratio working

---
 gusto/spatial_methods/augmentation.py   | 153 ++++++------------------
 gusto/timestepping/split_timestepper.py |  19 ++-
 2 files changed, 49 insertions(+), 123 deletions(-)

diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index 740afe31d..1c577c877 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -255,7 +255,7 @@ class MeanMixingRatio(Augmentation):
     Args:
         domain (:class:`Domain`): The domain object.
         eqns (:class:`PrognosticEquationSet`): The overarching equation set.
-        mixing_ratio (:class: list): List of mixing ratios that 
+        mixing_ratio (:class: list): A list of mixing ratios that 
         are to have augmented mean mixing ratio fields.
         OR, keep as a single mixing ratio, but define 
         multiple augmentations?
@@ -267,15 +267,16 @@ def __init__(
     ):
 
         self.name = 'mean_mixing_ratio'
+        
         self.mX_name = mX_name
         self.mean_name = 'mean_'+mX_name
-        print(self.mean_name)
 
         self.eqn_orig = eqns
         self.domain=domain
         exist_spaces = eqns.spaces
-
         self.idx_orig = len(exist_spaces)
+
+        # Change this to adjust for a list
         mean_idx = self.idx_orig
         self.mean_idx = mean_idx
 
@@ -285,12 +286,14 @@ def __init__(
 
         # Define the mean mixing ratio on the DG0 space
         DG0 = FunctionSpace(domain.mesh, "DG", 0)
-        self.DG0 = Function(DG0)
-        mX_space = eqns.spaces[mean_idx]
+        mX_space = eqns.spaces[mX_idx]
+
+        # But, what if the mixing ratios are in 
+        # different function spaces???
 
         self.mX_in = Function(mX_space)
-        self.mean_in = Function(mean_space)
-        self.compute_mean = Projector(mX_space, DG0)
+        self.mean = Function(DG0)
+        self.compute_mean = Projector(self.mX_in, self.mean)
 
          # Set up the scheme for the mean mixing ratio
 
@@ -302,101 +305,13 @@ def __init__(
         self.X = Function(self.fs)
         self.tests = TestFunctions(self.fs)
 
-        print(self.X)
-
         self.bcs = None
 
         self.x_in = Function(self.fs)
         self.x_out = Function(self.fs)
 
-        # Compute the new mean mass weighted term,
-        # IF this is conservatively transported.
-
-        #old_residual = eqns.residual
-
-        #mean_residual = old_residual.label_map(
-        #    lambda t: t.get(prognostic) == mX_name,
-        #    map_if_false=drop
-        #)
-        #mean_residual = prognostic.update_value(mean_residual, self.mean_name)
-
-        # Replace trial functions with those in the new mixed function space
-        #for term in eqns.residual:
-        #    print('\n')
-        #    print(term.form)
-
-        # If I do this later on with the transport terms, maybe I don't need to do this here?
-
-        #for idx in range(self.idx_orig):
-        #    field = eqns.field_names[idx]
-            # Seperate logic if mass-weighted or not?
-        #    print(idx)
-        #    print(field)
-
-        #    prog = split(self.X)[idx]
-
-        #    print('\n residual term before change')
-        #    print(old_residual.label_map(
-        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-        #        map_if_false=drop
-        #    ).form)
-
-        #    old_residual = old_residual.label_map(
-        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-        #        map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
-        #    )
-        #    old_residual = old_residual.label_map(
-        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-        #        map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
-        #    )
-            
-        #    print('\n residual term after change')
-         #   print(old_residual.label_map(
-        #        lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-        #        map_if_false=drop
-        #    ).form)
-
-            #old_residual = old_residual.label_map(
-            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-            #    map_if_true=replace_subject(self.X[idx], old_idx=idx)
-            #)
-            #old_residual = old_residual.label_map(
-            #    lambda t: t.get(prognostic) == field and not t.has_label(mass_weighted),
-            #    map_if_true=replace_test_function(self.tests[idx], old_idx=idx)
-            #)
-    
-        # Define the mean mixing ratio residual
-        #mean_residual = mX_residual.label_map(
-        #    lambda t: t.has_label(mass_weighted),
-            #map_if_true=replace_subject(mean_mass, old_idx=mX_idx),
-        #    map_if_false=replace_subject(self.X[mean_idx], old_idx = mX_idx)
-        #)
-
-        #print('\n mean mX residual before change')
-        #print(mean_residual.form)
-        #mean_residual = mean_residual.label_map(
-        #    all_terms,
-        #    replace_subject(self.X, old_idx=mX_idx, new_idx = mean_idx)
-        #)
-        #mean_residual = mean_residual.label_map(
-        #    all_terms,
-        #    replace_test_function(self.tests, old_idx=mX_idx, new_idx = mean_idx)
-        #)
-        #print('\n mean mX residual after change')
-        #print(mean_residual.form)
 
-        # Form the new residual
-        #self.residual = old_residual + mean_residual
-
-        #print('\n full residual')
-        #print(self.residual.form)
-
-        #print('yoyoyoy')
-
-        #for term in self.residual:
-        #    print(term.get(prognostic))
-
-    def setup_transport(self, spatial_methods, equation):
+    def setup_residual(self, spatial_methods, equation):
         # Copy spatial method for the mixing ratio onto the 
         # mean mixing ratio.
 
@@ -416,16 +331,16 @@ def setup_transport(self, spatial_methods, equation):
         for idx in range(self.idx_orig):
             field = self.eqn_orig.field_names[idx]
             # Seperate logic if mass-weighted or not?
-            print('\n', idx)
-            print(field)
+            #print('\n', idx)
+            #print(field)
 
             prog = split(self.X)[idx]
 
-            print('\n residual term before change')
-            print(old_residual.label_map(
-                lambda t: t.get(prognostic) == field,
-                map_if_false=drop
-            ).form)
+            #print('\n residual term before change')
+            #print(old_residual.label_map(
+            #    lambda t: t.get(prognostic) == field,
+            #    map_if_false=drop
+            #).form)
 
             old_residual = old_residual.label_map(
                 lambda t: t.get(prognostic) == field,
@@ -435,17 +350,17 @@ def setup_transport(self, spatial_methods, equation):
                 lambda t: t.get(prognostic) == field,
                 map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
             )
-            print('\n residual term after change')
-            print(old_residual.label_map(
-                lambda t: t.get(prognostic) == field,
-                map_if_false=drop
-            ).form)
+            #print('\n residual term after change')
+            #print(old_residual.label_map(
+            #    lambda t: t.get(prognostic) == field,
+            #    map_if_false=drop
+            #).form)
 
-        print('\n now setting up mean mixing ratio residual terms')
+        #print('\n now setting up mean mixing ratio residual terms')
 
 
-        print('\n mean mX residual after change')
-        print(mean_residual.form)
+        #print('\n mean mX residual after change')
+        #print(mean_residual.form)
         mean_residual = mean_residual.label_map(
             all_terms,
             replace_subject(self.X, old_idx=self.mX_idx, new_idx=self.mean_idx)
@@ -454,15 +369,17 @@ def setup_transport(self, spatial_methods, equation):
             all_terms,
             replace_test_function(self.tests, old_idx=self.mX_idx, new_idx=self.mean_idx)
         )
-        print('\n mean mX residual after change')
-        print(mean_residual.form)
+        #print('\n mean mX residual after change')
+        #print(mean_residual.form)
 
         # Form the new residual
-        self.residual = old_residual + mean_residual
+        residual = old_residual + mean_residual
+        self.residual = subject(residual, self.X)
 
         #Check these two forms
-        print('\n Original equation with residual of length, ', len(equation.residual))
-        print('\n Augmented equation with residual of length, ', len(self.residual))
+        #print('\n Original equation with residual of length, ', len(equation.residual))
+        #print('\n Augmented equation with residual of length, ', len(self.residual))
+
 
 
 
@@ -529,10 +446,10 @@ def update(self, x_in_mixed):
 
         # Compute the mean mixing ratio
         # How do I do this?
-        self.mX_in.assign(x_in_mixed[self.mX_idx])
+        self.mX_in.assign(x_in_mixed.subfunctions[self.mX_idx])
 
         # Project this into the lowest order space:
         self.compute_mean.project()
 
-        self.x_in.subfunctions[self.mean_idx].assign(self.mean_out)
+        self.x_in.subfunctions[self.mean_idx].assign(self.mean)
         pass
\ No newline at end of file
diff --git a/gusto/timestepping/split_timestepper.py b/gusto/timestepping/split_timestepper.py
index 98b9cb841..4738688d1 100644
--- a/gusto/timestepping/split_timestepper.py
+++ b/gusto/timestepping/split_timestepper.py
@@ -4,7 +4,7 @@
 from firedrake.fml import Label, drop
 from pyop2.profiling import timed_stage
 from gusto.core import TimeLevelFields, StateFields
-from gusto.core.labels import time_derivative, physics_label
+from gusto.core.labels import time_derivative, physics_label, dynamics_label
 from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation
 from gusto.timestepping.timestepper import BaseTimestepper, Timestepper
 from numpy import ones
@@ -299,13 +299,21 @@ def setup_scheme(self):
         print('Setting up base equation')
         self.setup_equation(self.equation)
 
-        # If there is an augmentation, set up these transport terms
-        # Or, perhaps set up the whole residual now ... ?
+         # If there is an augmentation, set up the residual now
         if self.scheme.augmentation is not None:
             if self.scheme.augmentation.name == 'mean_mixing_ratio':
-                self.scheme.augmentation.setup_transport(self.spatial_methods, self.equation)
+                self.scheme.augmentation.setup_residual(self.spatial_methods, self.equation)
                 print('Setting up augmented equation')
-
+                # Go through and label all non-physics terms with a "dynamics" label
+                dynamics = Label('dynamics')
+                self.scheme.augmentation.residual = self.scheme.augmentation.residual.label_map(
+                    lambda t: not any(t.has_label(time_derivative, physics_label)),
+                    map_if_true=lambda t: dynamics(t)
+                )
+                print(len(self.scheme.augmentation.residual.label_map(
+                    lambda t: t.has_label(dynamics),
+                    map_if_false=drop
+                )))
 
         # Go through and label all non-physics terms with a "dynamics" label
         dynamics = Label('dynamics')
@@ -316,6 +324,7 @@ def setup_scheme(self):
         if self.io.output.log_courant:
             self.scheme.courant_max = self.io.courant_max
 
+
     def setup_prescribed_expr(self, expr_func):
         """
         Sets up the prescribed transporting velocity, through a python function

From 965e15af9f0502ed644aac4f95c5946c87e85384 Mon Sep 17 00:00:00 2001
From: Tim Andrews <ta440@github.com>
Date: Wed, 15 Jan 2025 20:14:23 +0000
Subject: [PATCH 7/7] added a loopy kernel and limiter for the mean mixing
 ratio

---
 gusto/core/kernels.py                         |  48 +++++
 gusto/spatial_methods/augmentation.py         | 173 +++++++++++-------
 gusto/spatial_methods/limiters.py             |  78 +++++++-
 .../explicit_runge_kutta.py                   |   2 +
 4 files changed, 235 insertions(+), 66 deletions(-)

diff --git a/gusto/core/kernels.py b/gusto/core/kernels.py
index 6b78691dc..422b7456b 100644
--- a/gusto/core/kernels.py
+++ b/gusto/core/kernels.py
@@ -113,6 +113,54 @@ def apply(self, field, field_in):
                  {"field": (field, WRITE),
                   "field_in": (field_in, READ)})
 
+class MeanMixingRatioWeights():
+    """
+    Finds the lambda values for blending a mixing ratio and its
+    mean DG0 field in the MeanMixingRatioLimiter.
+
+    First, each cell is looped over and the minimum value is computed
+    """
+
+    def __init__(self, V):
+        """
+        Args:
+            V (:class:`FunctionSpace`): The space of the field to be clipped.
+        """
+        # Using DG1-equispaced, with 4 DOFs per cell
+        shapes = {'nDOFs': V.finat_element.space_dimension(),
+                  'nDOFs_base': int(V.finat_element.space_dimension() / 4)}
+        domain = "{{[i]: 0 <= i < {nDOFs_base}}}".format(**shapes)
+
+        instrs = ("""
+                  <float64> min_value = 0.0
+                  for i
+                      min_value = fmin(mX_field[i*4], mX_field[i*4+1])
+                      min_value = fmin(min_value, mX_field[i*4+2])
+                      min_value = fmin(min_value, mX_field[i*4+3])
+                      if min_value < 0.0
+                          lamda[i] = -min_value/(mean_field[i] - min_value)
+                      else
+                          lamda[i] = 0.0
+                      end
+                  end
+                  """)
+
+        self._kernel = (domain, instrs)
+
+    def apply(self, lamda, mX_field, mean_field):
+        """
+        Performs the par loop.
+
+        Args:
+            w (:class:`Function`): the field in which to store the weights. This
+                lives in the continuous target space.
+        """
+        par_loop(self._kernel, dx,
+                 {"lamda": (lamda, WRITE),
+                  "mX_field": (mX_field, READ),
+                  "mean_field": (mean_field, READ)})
+
+
 
 class MinKernel():
     """Finds the minimum DoF value of a field."""
diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py
index 1c577c877..e271061f0 100644
--- a/gusto/spatial_methods/augmentation.py
+++ b/gusto/spatial_methods/augmentation.py
@@ -20,9 +20,9 @@
     time_derivative, transport, transporting_velocity, TransportEquationType,
     logger, prognostic, mass_weighted
 )
-from gusto.spatial_methods.transport_methods import DGUpwind
+from gusto.spatial_methods.limiters import MeanLimiter, MixedFSLimiter
 import copy
-
+import numpy as np
 
 class Augmentation(object, metaclass=ABCMeta):
     """
@@ -55,6 +55,13 @@ def update(self, x_in_mixed):
 
         pass
 
+    def limit(self, x_in_mixed):
+        """
+        Apply any special limiting as part of the augmentation
+        """
+
+        pass
+
 
 class VorticityTransport(Augmentation):
     """
@@ -255,52 +262,64 @@ class MeanMixingRatio(Augmentation):
     Args:
         domain (:class:`Domain`): The domain object.
         eqns (:class:`PrognosticEquationSet`): The overarching equation set.
-        mixing_ratio (:class: list): A list of mixing ratios that 
-        are to have augmented mean mixing ratio fields.
-        OR, keep as a single mixing ratio, but define 
-        multiple augmentations?
-        START with just a single.
+        mX_names (:class: list): A list of mixing ratios that 
+        require augmented mean mixing ratio fields.
     """
 
     def __init__(
-            self, domain, eqns, mX_name
+            self, domain, eqns, mX_names
     ):
 
         self.name = 'mean_mixing_ratio'
-        
-        self.mX_name = mX_name
-        self.mean_name = 'mean_'+mX_name
+        self.mX_names = mX_names
+        self.mX_num = len(mX_names)
 
+        # Store information about original equation set
         self.eqn_orig = eqns
-        self.domain=domain
+        self.domain = domain
         exist_spaces = eqns.spaces
         self.idx_orig = len(exist_spaces)
 
-        # Change this to adjust for a list
-        mean_idx = self.idx_orig
-        self.mean_idx = mean_idx
-
-        # Extract the mixing ratio in question:
-        mX_idx = eqns.field_names.index(mX_name)
-        self.mX_idx = mX_idx
-
         # Define the mean mixing ratio on the DG0 space
         DG0 = FunctionSpace(domain.mesh, "DG", 0)
-        mX_space = eqns.spaces[mX_idx]
-
-        # But, what if the mixing ratios are in 
-        # different function spaces???
-
-        self.mX_in = Function(mX_space)
-        self.mean = Function(DG0)
-        self.compute_mean = Projector(self.mX_in, self.mean)
-
-         # Set up the scheme for the mean mixing ratio
-
-        mean_mX = Function(DG0, name=self.mean_name)
-        mean_space = DG0
-        exist_spaces.append(mean_space)
 
+        # Set up fields and names for each mixing ratio
+        self.mean_names = []
+        self.mean_idxs = []
+        self.mX_idxs = []
+        mX_spaces = []
+        mean_spaces = []
+        self.limiters = []
+        #self.sublimiters = {}
+
+        for i in range(self.mX_num):
+            mX_name = mX_names[i]
+            print(mX_names)
+            print(mX_name)
+            self.mean_names.append('mean_'+mX_name)
+            mean_spaces.append(DG0)
+            exist_spaces.append(DG0)
+            self.mean_idxs.append(self.idx_orig + i)
+
+            # Extract the mixing ratio in question:
+            mX_idx = eqns.field_names.index(mX_name)
+            mX_spaces.append(eqns.spaces[mX_idx])
+            self.mX_idxs.append(mX_idx)
+
+            # Define a limiter
+            self.limiters.append(MeanLimiter(eqns.spaces[mX_idx]))
+            #self.sublimiters.update({mX_name: MeanLimiter(eqns.spaces[mX_idx])})
+
+        #self.limiter = MixedFSLimiter(self.eqn_orig, self.sublimiters)
+        #self.limiter = MixedFSLimiter(sublimiters)
+
+        # Contruct projectors for computing the mean mixing ratios
+        self.mX_ins = [Function(mX_spaces[i]) for i in range(self.mX_num)]
+        self.mean_outs = [Function(mean_spaces[i]) for i in range(self.mX_num)]
+        self.compute_means = [Projector(self.mX_ins[i], self.mean_outs[i]) \
+                               for i in range(self.mX_num)]
+
+        # Create the new mixed function space:
         self.fs = MixedFunctionSpace(exist_spaces)
         self.X = Function(self.fs)
         self.tests = TestFunctions(self.fs)
@@ -315,14 +334,24 @@ def setup_residual(self, spatial_methods, equation):
         # Copy spatial method for the mixing ratio onto the 
         # mean mixing ratio.
 
-        old_residual = equation.residual
+        orig_residual = equation.residual
 
         # Copy the mean mixing ratio residual terms:
-        mean_residual = old_residual.label_map(
-            lambda t: t.get(prognostic) == self.mX_name,
-            map_if_false=drop
-        )
-        mean_residual = prognostic.update_value(mean_residual, self.mean_name)
+        for i in range(self.mX_num):
+            if i == 0:
+                mean_residual = orig_residual.label_map(
+                    lambda t: t.get(prognostic) == self.mX_names[i],
+                    map_if_false=drop
+                )
+                mean_residual = prognostic.update_value(mean_residual, self.mean_names[i])
+            else:
+                mean_residual_term = orig_residual.label_map(
+                    lambda t: t.get(prognostic) == self.mX_names[i],
+                    map_if_false=drop
+                )
+                mean_residual_term = prognostic.update_value(mean_residual_term,\
+                                                              self.mean_names[i])
+                mean_residual = mean_residual + mean_residual_term
 
         print('\n in setup_transport')
 
@@ -334,7 +363,7 @@ def setup_residual(self, spatial_methods, equation):
             #print('\n', idx)
             #print(field)
 
-            prog = split(self.X)[idx]
+            #prog = split(self.X)[idx]
 
             #print('\n residual term before change')
             #print(old_residual.label_map(
@@ -342,11 +371,11 @@ def setup_residual(self, spatial_methods, equation):
             #    map_if_false=drop
             #).form)
 
-            old_residual = old_residual.label_map(
+            orig_residual = orig_residual.label_map(
                 lambda t: t.get(prognostic) == field,
                 map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
             )
-            old_residual = old_residual.label_map(
+            orig_residual = orig_residual.label_map(
                 lambda t: t.get(prognostic) == field,
                 map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
             )
@@ -361,19 +390,25 @@ def setup_residual(self, spatial_methods, equation):
 
         #print('\n mean mX residual after change')
         #print(mean_residual.form)
-        mean_residual = mean_residual.label_map(
-            all_terms,
-            replace_subject(self.X, old_idx=self.mX_idx, new_idx=self.mean_idx)
-        )
-        mean_residual = mean_residual.label_map(
-            all_terms,
-            replace_test_function(self.tests, old_idx=self.mX_idx, new_idx=self.mean_idx)
-        )
+
+        # Update the subject and test functions for the
+        # mean mixing ratios
+        for i in range(self.mX_num):
+            mean_residual = mean_residual.label_map(
+                all_terms,
+                replace_subject(self.X, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i])
+            )
+            mean_residual = mean_residual.label_map(
+                all_terms,
+                replace_test_function(self.tests, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i])
+            )
+
+
         #print('\n mean mX residual after change')
         #print(mean_residual.form)
 
         # Form the new residual
-        residual = old_residual + mean_residual
+        residual = orig_residual + mean_residual
         self.residual = subject(residual, self.X)
 
         #Check these two forms
@@ -383,6 +418,7 @@ def setup_residual(self, spatial_methods, equation):
 
 
 
+
     def setup_transport_old(self, spatial_methods):
         mX_spatial_method = next(method for method in spatial_methods if method.variable == self.mX_name)
         
@@ -427,6 +463,7 @@ def pre_apply(self, x_in):
 
     def post_apply(self, x_out):
         """
+        Apply the limiters.
         Sets the output fields, i.e. not the mean mixing ratios
 
         Args:
@@ -438,18 +475,28 @@ def post_apply(self, x_out):
 
     def update(self, x_in_mixed):
         """
-        ...
+        Compute the mean mixing ratio field by projecting the mixing 
+        ratio from DG1 into DG0.
+
+        SHouldn't this be a conservative projection??!!
+        This requires density fields...
 
         Args:
             x_in_mixed (:class:`Function`): The mixed function to update.
         """
-
-        # Compute the mean mixing ratio
-        # How do I do this?
-        self.mX_in.assign(x_in_mixed.subfunctions[self.mX_idx])
-
-        # Project this into the lowest order space:
-        self.compute_mean.project()
-
-        self.x_in.subfunctions[self.mean_idx].assign(self.mean)
-        pass
\ No newline at end of file
+        print('in update')
+        for i in range(self.mX_num):
+            self.mX_ins[i].assign(x_in_mixed.subfunctions[self.mX_idxs[i]])
+            self.compute_means[i].project()
+            self.x_in.subfunctions[self.mean_idxs[i]].assign(self.mean_outs[i])
+
+    def limit(self, x_in_mixed):
+        # Ensure non-negativity by applying the blended limiter
+        for i in range(self.mX_num):
+            print('limiting within the augmentation')
+            mX_field = x_in_mixed.subfunctions[self.mX_idxs[i]]
+            mean_field = x_in_mixed.subfunctions[self.mean_idxs[i]]
+            print(np.min(x_in_mixed.subfunctions[self.mX_idxs[i]].dat.data))
+            self.limiters[i].apply(mX_field, mean_field)
+            print(np.min(x_in_mixed.subfunctions[self.mX_idxs[i]].dat.data))
+            #x_in_mixed.subfunctions[self.mX_idxs[i]].assign(mX_field)
diff --git a/gusto/spatial_methods/limiters.py b/gusto/spatial_methods/limiters.py
index 8a782a054..c04311540 100644
--- a/gusto/spatial_methods/limiters.py
+++ b/gusto/spatial_methods/limiters.py
@@ -6,14 +6,14 @@
 """
 
 from firedrake import (BrokenElement, Function, FunctionSpace, interval,
-                       FiniteElement, TensorProductElement)
+                       FiniteElement, TensorProductElement, Constant)
 from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter
-from gusto.core.kernels import LimitMidpoints, ClipZero
+from gusto.core.kernels import LimitMidpoints, ClipZero, MeanMixingRatioWeights
 
 import numpy as np
 
 __all__ = ["DG1Limiter", "ThetaLimiter", "NoLimiter", "ZeroLimiter",
-           "MixedFSLimiter"]
+           "MixedFSLimiter", "MeanLimiter"]
 
 
 class DG1Limiter(object):
@@ -261,3 +261,75 @@ def apply(self, fields):
         for field, sublimiter in self.sublimiters.items():
             field = fields.subfunctions[self.field_idxs[field]]
             sublimiter.apply(field)
+
+class MeanLimiter(object):
+    """
+    A limiter for a mixing ratio that ensures non-negativity by blending
+    the mixing ratio field with that of a mean mixing ratio. The blending
+    factor is given by the DG0 function lamda.
+    """
+
+    def __init__(self, space):
+        """
+        Args:
+            space: The mixed function space for the equation set
+            mX_field: The mixing ratio field
+            mean_field: The mean mixing ratio field
+        Raises:
+            ValueError: If the space is not appropriate for the limiter.
+        """
+
+        self.space = space 
+        mesh = space.mesh()
+
+        # check that space is DG1
+        degree = space.ufl_element().degree()
+        if (space.ufl_element().sobolev_space.name != 'L2'
+            or ((type(degree) is tuple and np.any([deg != 1 for deg in degree]))
+                and degree != 1)):
+            raise NotImplementedError('MeanLimiter only implemented for mixing' +
+                                      'ratios in the DG1 space')
+
+        # Create equispaced DG1 space needed for limiting
+        if space.extruded:
+            cell = mesh._base_mesh.ufl_cell().cellname()
+            DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced")
+            DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced")
+            DG1_element = TensorProductElement(DG1_hori_elt, DG1_vert_elt)
+        else:
+            cell = mesh.ufl_cell().cellname()
+            DG1_element = FiniteElement("DG", cell, 1, variant="equispaced")
+
+        DG1_equispaced = FunctionSpace(mesh, DG1_element)
+        print(DG1_equispaced.finat_element.space_dimension())
+
+        DG0 = FunctionSpace(mesh, 'DG', 0)
+
+        self.lamda = Function(DG1_equispaced)
+        self.mX_field = Function(DG1_equispaced)
+        self.mean_field = Function(DG0)
+
+        self._kernel = MeanMixingRatioWeights(self.space)
+
+    def apply(self, mX_field, mean_field):
+        """
+        The application of the limiter to the field.
+
+        Args:
+            field (:class:`Function`): the field to apply the limiter to.
+
+        Raises:
+             AssertionError: If the field is not in the correct space.
+         """
+
+        #self.field_old.interpolate(mX_field)
+        #self.mean_field.interpolate(mean_field)
+
+        # Compute the weights, lamda:
+        self._kernel.apply(self.lamda, mX_field, mean_field)
+
+        print('applying limiter')
+        #print(self.lamda.dat.data)
+
+        # Compute the blended field
+        mX_field.interpolate((Constant(1.0) - self.lamda)*mX_field + self.lamda*mean_field)
diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py
index 16732bd50..bbd37e565 100644
--- a/gusto/time_discretisation/explicit_runge_kutta.py
+++ b/gusto/time_discretisation/explicit_runge_kutta.py
@@ -477,6 +477,8 @@ def apply_cycle(self, x_out, x_in):
 
         for i in range(self.nStages):
             self.solve_stage(x_in, i)
+            if self.augmentation is not None:
+                self.augmentation.limit(self.x1)
         x_out.assign(self.x1)