Coupled transport

gusto/active_tracers.py

@@ -78,8 +78,6 @@ def __init__(self, name, space, variable_type,
         self.variable_type = variable_type
         self.phase = phase
         self.chemical = chemical
-        if self.variable_type != TracerVariableType.mixing_ratio:
-            raise NotImplementedError('Only mixing ratio tracers are currently implemented')
 
         if (variable_type == TracerVariableType.density and transport_eqn == TransportEquationType.advective):
             logger.warning('Active tracer initialised which describes a '

gusto/equations.py

@@ -80,7 +80,6 @@ def __init__(self, domain, function_space, field_name, Vu=None):
             Vu (:class:`FunctionSpace`, optional): the function space for the
                 velocity field. If this is not specified, uses the HDiv spaces
                 set up by the domain. Defaults to None.
-            **kwargs: any keyword arguments to be passed to the advection form.
         """
         super().__init__(domain, function_space, field_name)
 
@@ -508,36 +507,38 @@ def get_active_tracer(self, field_name):
         return active_tracer_to_return
 
 
-class ForcedAdvectionEquation(PrognosticEquationSet):
+class CoupledTransportEquation(PrognosticEquationSet):
     u"""
-    Discretises the advection equation with a source/sink term,               \n
+    Discretises the transport equation,               \n
     ∂q/∂t + (u.∇)q = F,
-    which can also be augmented with active tracers.
+    with the application of active tracers.
+    As there are multiple tracers or species that are
+    interacting, q and F are vectors.
+    This equation can be enhanced through the addition of
+    sources or sinks (F) by applying it with physics schemes.
     """
-    def __init__(self, domain, function_space, field_name, Vu=None,
-                 active_tracers=None, **kwargs):
+    def __init__(self, domain, active_tracers, Vu=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
                 mesh and the compatible function spaces.
-            function_space (:class:`FunctionSpace`): the function space that the
-                equation's prognostic is defined on.
-            field_name (str): name of the prognostic field.
+            active_tracers (list): a list of `ActiveTracer` objects
+                that encode the metadata for any active tracers to be included
+                in the equations. This is required for using this class; if there
+                is only a field to be advected, use the AdvectionEquation
+                instead.
             Vu (:class:`FunctionSpace`, optional): the function space for the
                 velocity field. If this is not specified, uses the HDiv spaces
                 set up by the domain. Defaults to None.
-            active_tracers (list, optional): a list of `ActiveTracer` objects
-                that encode the metadata for any active tracers to be included
-                in the equations. Defaults to None.
         """
 
-        self.field_names = [field_name]
-        self.space_names = {field_name: function_space.name}
         self.active_tracers = active_tracers
         self.terms_to_linearise = {}
+        self.field_names = []
+        self.space_names = {}
 
         # Build finite element spaces
-        self.spaces = [domain.spaces("tracer", V=function_space)]
+        self.spaces = []
 
         # Add active tracers to the list of prognostics
         if active_tracers is None:
@@ -547,27 +548,25 @@ def __init__(self, domain, function_space, field_name, Vu=None,
         # Make the full mixed function space
         W = MixedFunctionSpace(self.spaces)
 
-        # Can now call the underlying PrognosticEquation
         full_field_name = "_".join(self.field_names)
         PrognosticEquation.__init__(self, domain, W, full_field_name)
 
         if Vu is not None:
             V = domain.spaces("HDiv", V=Vu, overwrite_space=True)
         else:
             V = domain.spaces("HDiv")
-        u = self.prescribed_fields("u", V)
+        _ = self.prescribed_fields("u", V)
 
         self.tests = TestFunctions(W)
         self.X = Function(W)
 
         mass_form = self.generate_mass_terms()
-        transport_form = prognostic(advection_form(self.tests[0], split(self.X)[0], u), field_name)
 
-        self.residual = subject(mass_form + transport_form, self.X)
+        self.residual = subject(mass_form, self.X)
 
         # Add transport of tracers
-        if len(active_tracers) > 0:
-            self.residual += self.generate_tracer_transport_terms(active_tracers)
+        self.residual += self.generate_tracer_transport_terms(active_tracers)
+
 
 # ============================================================================ #
 # Specified Equation Sets

integration-tests/equations/test_coupled_transport.py

@@ -0,0 +1,65 @@
+"""
+Tests the CoupledTransportEquation class.
+Two tracers are transported using combinations of advective and
+conservative forms. The tracers are set to be a mixing ratio when
+using the advective form and a density for the conservative form.
+
+"""
+
+from firedrake import norm
+from gusto import *
+import pytest
+
+
+def run(timestepper, tmax, f_end):
+    timestepper.run(0, tmax)
+    norm1 = norm(timestepper.fields("f1") - f_end) / norm(f_end)
+    norm2 = norm(timestepper.fields("f2") - f_end) / norm(f_end)
+    return norm1, norm2
+
+
+@pytest.mark.parametrize("geometry", ["slice", "sphere"])
+@pytest.mark.parametrize("equation_form1", ["advective", "continuity"])
+@pytest.mark.parametrize("equation_form2", ["advective", "continuity"])
+def test_coupled_transport_scalar(tmpdir, geometry, equation_form1, equation_form2, tracer_setup):
+    setup = tracer_setup(tmpdir, geometry)
+    domain = setup.domain
+
+    if equation_form1 == "advective":
+        tracer1 = ActiveTracer(name='f1', space='DG',
+                               variable_type=TracerVariableType.mixing_ratio,
+                               transport_eqn=TransportEquationType.advective)
+    else:
+        tracer1 = ActiveTracer(name='f1', space='DG',
+                               variable_type=TracerVariableType.density,
+                               transport_eqn=TransportEquationType.conservative)
+
+    if equation_form2 == "advective":
+        tracer2 = ActiveTracer(name='f2', space='DG',
+                               variable_type=TracerVariableType.mixing_ratio,
+                               transport_eqn=TransportEquationType.advective)
+    else:
+        tracer2 = ActiveTracer(name='f2', space='DG',
+                               variable_type=TracerVariableType.density,
+                               transport_eqn=TransportEquationType.conservative)
+
+    tracers = [tracer1, tracer2]
+
+    V = domain.spaces("HDiv")
+    eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V)
+
+    transport_scheme = SSPRK3(domain)
+    transport_method = [DGUpwind(eqn, 'f1'), DGUpwind(eqn, 'f2')]
+
+    timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method)
+
+    # Initial conditions
+    timestepper.fields("f1").interpolate(setup.f_init)
+    timestepper.fields("f2").interpolate(setup.f_init)
+    timestepper.fields("u").project(setup.uexpr)
+
+    error1, error2 = run(timestepper, setup.tmax, setup.f_end)
+    assert error1 < setup.tol, \
+        'The transport error for tracer 1 is greater than the permitted tolerance'
+    assert error2 < setup.tol, \
+        'The transport error for tracer 2 is greater than the permitted tolerance'

integration-tests/equations/test_forced_advection.py

@@ -45,10 +45,15 @@ def run_forced_advection(tmpdir):
     msat = Function(VD)
     msat.interpolate(msat_expr)
 
-    rain = Rain(space='tracer',
+    # Rain is a first tracer
+    rain = Rain(space='DG',
                 transport_eqn=TransportEquationType.no_transport)
-    meqn = ForcedAdvectionEquation(domain, VD, field_name="water_vapour", Vu=Vu,
-                                   active_tracers=[rain])
+
+    # Also, have water_vapour as a tracer:
+    water_vapour = WaterVapour(space='DG')
+
+    meqn = CoupledTransportEquation(domain, active_tracers=[rain, water_vapour], Vu=Vu)
+
     transport_method = DGUpwind(meqn, "water_vapour")
     physics_parametrisations = [InstantRain(meqn, msat, rain_name="rain",
                                             parameters=None)]
@@ -72,6 +77,10 @@ def run_forced_advection(tmpdir):
     stepper.fields("u").project(as_vector([u_max]))
     stepper.fields("water_vapour").project(mexpr)
 
+    # Start with no rain:
+    no_rain = 0*x
+    stepper.fields("rain").interpolate(no_rain)
+
     # exact rainfall profile (analytically)
     r_exact = stepper.fields("r_exact", space=VD)
     lim1 = Lx/(2*pi) * acos((C0 + K0 - Csat)/Ksat)

integration-tests/physics/test_precipitation.py

@@ -30,11 +30,15 @@ def setup_fallout(dirname):
     domain = Domain(mesh, dt, "CG", 1)
     x = SpatialCoordinate(mesh)
 
+    # Define the tracers
+    rho = ActiveTracer(name='rho', space='DG1_equispaced', variable_type=TracerVariableType.density)
+    rain = Rain(space='DG1_equispaced')
+
+    Vu = domain.spaces("HDiv")
+
     # Equation
-    Vrho = domain.spaces("DG1_equispaced")
-    active_tracers = [Rain(space='DG1_equispaced')]
-    eqn = ForcedAdvectionEquation(domain, Vrho, "rho", active_tracers=active_tracers)
-    transport_method = DGUpwind(eqn, "rho")
+    eqn = CoupledTransportEquation(domain, active_tracers=[rho, rain], Vu=Vu)
+    transport_method = [DGUpwind(eqn, "rho"), DGUpwind(eqn, "rain")]
 
     # I/O
     output = OutputParameters(dirname=dirname+"/fallout", dumpfreq=10, dumplist=['rain'])