PR #469: from firedrakeproject/tracer_conservative_transport

Add conservative form of tracer transport equation, and an expansion of the RK discretisations to work with this

gusto/active_tracers.py

@@ -53,6 +53,7 @@ class ActiveTracer(object):
     """
     def __init__(self, name, space, variable_type,
                  transport_eqn=TransportEquationType.advective,
+                 density_name=None,
                  phase=Phases.gas, chemical=None):
         """
         Args:
@@ -63,6 +64,10 @@ def __init__(self, name, space, variable_type,
             transport_eqn (:class:`TransportEquationType`, optional): enumerator
                 indicating the type of transport equation to be solved (e.g.
                 advective). Defaults to `TransportEquationType.advective`.
+            density_name (str): the name of the associated density for a mixing
+                ratio when using the tracer_conservative transport. Defaults to None,
+                but raises an error if tracer_conservative transport is used without
+                a specified density.
             phase (:class:`Phases`, optional): enumerator indicating the phase
                 of the tracer variable. Defaults to `Phases.gas`.
             chemical (str, optional): string to describe the chemical that this
@@ -75,14 +80,21 @@ def __init__(self, name, space, variable_type,
         self.name = name
         self.space = space
         self.transport_eqn = transport_eqn
+        self.density_name = density_name
         self.variable_type = variable_type
         self.phase = phase
         self.chemical = chemical
 
-        if (variable_type == TracerVariableType.density and transport_eqn == TransportEquationType.advective):
+        if (variable_type == TracerVariableType.density
+                and transport_eqn == TransportEquationType.advective):
             logger.warning('Active tracer initialised which describes a '
                            + 'density but solving the advective transport eqn')
 
+        if (transport_eqn == TransportEquationType.tracer_conservative
+                and density_name is None):
+            raise ValueError(f'Active tracer {name} using tracer conservative '
+                             + 'transport needs an associated density.')
+
 
 class WaterVapour(ActiveTracer):
     """An object encoding the details of water vapour as a tracer."""

gusto/common_forms.py

@@ -12,7 +12,7 @@
 __all__ = ["advection_form", "continuity_form", "vector_invariant_form",
            "kinetic_energy_form", "advection_equation_circulation_form",
            "diffusion_form", "linear_advection_form", "linear_continuity_form",
-           "split_continuity_form"]
+           "split_continuity_form", "tracer_conservative_form"]
 
 
 def advection_form(test, q, ubar):
@@ -262,3 +262,28 @@ def split_continuity_form(equation):
                 map_if_true=drop)
 
     return equation
+
+
+def tracer_conservative_form(test, q, rho, ubar):
+    u"""
+    The form corresponding to the continuity transport operator.
+
+    This describes ∇.(u*q*rho) for transporting velocity u and a
+    transported tracer (mixing ratio), q, with an associated density, rho.
+
+    Args:
+        test (:class:`TestFunction`): the test function.
+        q (:class:`ufl.Expr`): the tracer to be transported.
+        rho (:class:`ufl.Expr`): the reference density that will
+        mulitply with q before taking the divergence.
+        ubar (:class:`ufl.Expr`): the transporting velocity.
+
+    Returns:
+        class:`LabelledForm`: a labelled transport form.
+    """
+
+    q_rho = q*rho
+    L = inner(test, div(outer(q_rho, ubar)))*dx
+    form = transporting_velocity(L, ubar)
+
+    return transport(form, TransportEquationType.tracer_conservative)

gusto/configuration.py

@@ -31,13 +31,16 @@ class TransportEquationType(Enum):
     conservative: ∂q/∂t + ∇.(u*q) = 0                                         \n
     vector_invariant: ∂q/∂t + (∇×q)×u + (1/2)∇(q.u) + (1/2)[(∇q).u -(∇u).q)] = 0
     circulation: ∂q/∂t + (∇×q)×u + non-transport terms = 0
+    tracer_conservative: ∂(q*rho)/∂t + ∇.(u*q*rho) = 0, for a reference density of rho
+    for the tracer, q.
     """
 
     no_transport = 702
     advective = 19
     conservative = 291
     vector_invariant = 9081
     circulation = 512
+    tracer_conservative = 296
 
 
 class Configuration(object):

gusto/equations.py

@@ -20,7 +20,8 @@
 from gusto.common_forms import (
     advection_form, continuity_form, vector_invariant_form,
     kinetic_energy_form, advection_equation_circulation_form,
-    diffusion_form, linear_continuity_form, linear_advection_form
+    diffusion_form, linear_continuity_form, linear_advection_form,
+    tracer_conservative_form
 )
 from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType
 from gusto.configuration import TransportEquationType
@@ -430,6 +431,43 @@ def add_tracers_to_prognostics(self, domain, active_tracers):
             else:
                 raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}')
 
+    def generate_tracer_mass_terms(self, active_tracers):
+        """
+        Adds the mass forms for the active tracers to the equation set.
+
+        Args:
+            active_tracers (list): A list of :class:`ActiveTracer` objects that
+                encode the metadata for the active tracers.
+
+        Returns:
+            :class:`LabelledForm`: a labelled form containing the mass
+                terms for the active tracers. This is the usual mass form
+                unless using tracer_conservative, where it is multiplied
+                by the reference density.
+        """
+
+        for i, tracer in enumerate(active_tracers):
+            idx = self.field_names.index(tracer.name)
+            tracer_prog = split(self.X)[idx]
+            tracer_test = self.tests[idx]
+
+            if tracer.transport_eqn == TransportEquationType.tracer_conservative:
+                ref_density_idx = self.field_names.index(tracer.density_name)
+                ref_density = split(self.X)[ref_density_idx]
+                q = tracer_prog*ref_density
+                mass = subject(prognostic(inner(q, tracer_test)*dx,
+                                          self.field_names[idx]), self.X)
+            else:
+                mass = subject(prognostic(inner(tracer_prog, tracer_test)*dx,
+                                          self.field_names[idx]), self.X)
+
+            if i == 0:
+                mass_form = time_derivative(mass)
+            else:
+                mass_form += time_derivative(mass)
+
+        return mass_form
+
     def generate_tracer_transport_terms(self, active_tracers):
         """
         Adds the transport forms for the active tracers to the equation set.
@@ -473,6 +511,13 @@ def generate_tracer_transport_terms(self, active_tracers):
                     tracer_adv = prognostic(
                         continuity_form(tracer_test, tracer_prog, u),
                         tracer.name)
+                elif tracer.transport_eqn == TransportEquationType.tracer_conservative:
+                    ref_density_idx = self.field_names.index(tracer.density_name)
+                    ref_density = split(self.X)[ref_density_idx]
+                    tracer_adv = prognostic(
+                        tracer_conservative_form(tracer_test, tracer_prog,
+                                                 ref_density, u), tracer.name)
+
                 else:
                     raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised')
 
@@ -562,9 +607,9 @@ def __init__(self, domain, active_tracers, Vu=None):
         self.tests = TestFunctions(W)
         self.X = Function(W)
 
-        mass_form = self.generate_mass_terms()
-
-        self.residual = subject(mass_form, self.X)
+        # Add mass forms for the tracers, which will use
+        # mass*density for any tracer_conservative terms
+        self.residual = self.generate_tracer_mass_terms(active_tracers)
 
         # Add transport of tracers
         self.residual += self.generate_tracer_transport_terms(active_tracers)
@@ -574,7 +619,6 @@ def __init__(self, domain, active_tracers, Vu=None):
 # Specified Equation Sets
 # ============================================================================ #
 
-
 class ShallowWaterEquations(PrognosticEquationSet):
     u"""
     Class for the (rotating) shallow-water equations, which evolve the velocity