PR #504: from firedrakeproject/conservative_physics

Conservative transport with physics

gusto/core/labels.py

@@ -95,6 +95,7 @@ def __call__(self, target, value=None):
 transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor, ufl.indexed.Indexed])
 prognostic = Label("prognostic", validator=lambda value: type(value) == str)
 linearisation = Label("linearisation", validator=lambda value: type(value) in [LabelledForm, Term])
+mass_weighted = Label("mass_weighted", validator=lambda value: type(value) in [LabelledForm, Term])
 ibp_label = Label("ibp", validator=lambda value: type(value) == IntegrateByParts)
 
 # labels for terms in the equations

gusto/equations/prognostic_equations.py

@@ -10,7 +10,7 @@
     replace_subject, replace_trial_function
 )
 from gusto.core import PrescribedFields
-from gusto.core.labels import time_derivative, prognostic, linearisation
+from gusto.core.labels import time_derivative, prognostic, linearisation, mass_weighted
 from gusto.equations.common_forms import (
     advection_form, continuity_form, tracer_conservative_form
 )
@@ -142,9 +142,31 @@ def generate_mass_terms(self):
             :class:`LabelledForm`: a labelled form containing the mass terms.
         """
 
+        if self.active_tracers is None:
+            tracer_names = []
+        else:
+            tracer_names = [tracer.name for tracer in self.active_tracers]
+
         for i, (test, field_name) in enumerate(zip(self.tests, self.field_names)):
             prog = split(self.X)[i]
             mass = subject(prognostic(inner(prog, test)*dx, field_name), self.X)
+
+            # Check if the field is a conservatively transported tracer. If so,
+            # create a mass-weighted mass form and store this and the original
+            # mass form in a mass-weighted label
+            for j, tracer_name in enumerate(tracer_names):
+                if field_name == tracer_name:
+                    if self.active_tracers[j].transport_eqn == TransportEquationType.tracer_conservative:
+                        standard_mass_form = mass
+
+                        # The mass-weighted mass form is multiplied by the reference density
+                        ref_density_idx = self.field_names.index(self.active_tracers[j].density_name)
+                        ref_density = split(self.X)[ref_density_idx]
+                        q = prog*ref_density
+                        mass_weighted_form = time_derivative(subject(prognostic(inner(q, test)*dx,
+                                                             field_name), self.X))
+
+                        mass = mass_weighted(standard_mass_form, mass_weighted_form)
             if i == 0:
                 mass_form = time_derivative(mass)
             else:
@@ -304,43 +326,6 @@ 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.
@@ -377,29 +362,35 @@ def generate_tracer_transport_terms(self, active_tracers):
                 tracer_prog = split(self.X)[idx]
                 tracer_test = self.tests[idx]
                 if tracer.transport_eqn == TransportEquationType.advective:
-                    tracer_adv = prognostic(
+                    tracer_adv = subject(prognostic(
                         advection_form(tracer_test, tracer_prog, u),
-                        tracer.name)
+                        tracer.name), self.X)
                 elif tracer.transport_eqn == TransportEquationType.conservative:
-                    tracer_adv = prognostic(
+                    tracer_adv = subject(prognostic(
                         continuity_form(tracer_test, tracer_prog, u),
-                        tracer.name)
+                        tracer.name), self.X)
                 elif tracer.transport_eqn == TransportEquationType.tracer_conservative:
+                    default_adv_form = subject(prognostic(
+                        advection_form(tracer_test, tracer_prog, u),
+                        tracer.name), self.X)
+
                     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)
+                    mass_weighted_tracer_adv = subject(prognostic(
+                        tracer_conservative_form(tracer_test, tracer_prog, ref_density, u),
+                        tracer.name), self.X)
 
+                    # Store the conservative transport form in the mass_weighted label,
+                    # but by default use an advective form.
+                    tracer_adv = mass_weighted(default_adv_form, mass_weighted_tracer_adv)
                 else:
                     raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised')
-
                 if no_tracer_transported:
                     # We arrive here for the first tracer to be transported
-                    adv_form = subject(tracer_adv, self.X)
+                    adv_form = tracer_adv
                     no_tracer_transported = False
                 else:
-                    adv_form += subject(tracer_adv, self.X)
+                    adv_form += tracer_adv
 
         return adv_form
 

gusto/equations/transport_equations.py

@@ -122,7 +122,7 @@ def __init__(self, domain, active_tracers, Vu=None):
 
         # 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)
+        self.residual = self.generate_mass_terms()
 
         # Add transport of tracers
         self.residual += self.generate_tracer_transport_terms(active_tracers)

gusto/spatial_methods/transport_methods.py

@@ -8,7 +8,8 @@
 )
 from firedrake.fml import Term, keep, drop
 from gusto.core.configuration import IntegrateByParts, TransportEquationType
-from gusto.core.labels import prognostic, transport, transporting_velocity, ibp_label
+from gusto.core.labels import (prognostic, transport, transporting_velocity, ibp_label,
+                               mass_weighted)
 from gusto.core.logging import logger
 from gusto.spatial_methods.spatial_methods import SpatialMethod
 
@@ -34,7 +35,12 @@ def __init__(self, equation, variable):
         # Inherited init method extracts original term to be replaced
         super().__init__(equation, variable, transport)
 
-        self.transport_equation_type = self.original_form.terms[0].get(transport)
+        # If this is term has a mass_weighted label, then we need to
+        # use the tracer_conservative version of the transport method.
+        if self.original_form.terms[0].has_label(mass_weighted):
+            self.transport_equation_type = TransportEquationType.tracer_conservative
+        else:
+            self.transport_equation_type = self.original_form.terms[0].get(transport)
 
         if self.transport_equation_type == TransportEquationType.tracer_conservative:
             # Extract associated density of the variable
@@ -77,6 +83,22 @@ def replace_form(self, equation):
             # Create new term
             new_term = Term(self.form.form, original_term.labels)
 
+            # Check if this is a conservative transport
+            if original_term.has_label(mass_weighted):
+                # Extract the original and discretised mass_weighted terms
+                original_mass_weighted_term = original_term.get(mass_weighted).terms[0]
+                new_mass_weighted = self.form.terms[0].get(mass_weighted)
+
+                # Ensure the correct labels for the new mass weighted term
+                new_mass_weighted_term = Term(new_mass_weighted.form, original_mass_weighted_term.labels)
+                # Update the mass weighted transporting velocity
+                new_mass_weighted_transporting_velocity = new_mass_weighted.terms[0].get(transporting_velocity)
+                new_mass_weighted_term = transporting_velocity.update_value(new_mass_weighted_term, new_mass_weighted_transporting_velocity)
+
+                # Add the discretised mass weighted transport term as the
+                # new mass weighted label.
+                new_term = mass_weighted.update_value(new_term, new_mass_weighted_term)
+
             # Replace original term with new term
             equation.residual = equation.residual.label_map(
                 lambda t: t.has_label(transport) and t.get(prognostic) == self.variable,
@@ -199,14 +221,20 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE,
                                                     self.field, ibp=ibp)
 
             elif self.transport_equation_type == TransportEquationType.tracer_conservative:
-                form = upwind_tracer_conservative_form(self.domain, self.test,
+                mass_weighted_form = upwind_tracer_conservative_form(self.domain, self.test,
+                                                                     self.field,
+                                                                     self.conservative_density,
+                                                                     ibp=ibp)
+                advective_form = upwind_advection_form(self.domain, self.test,
                                                        self.field,
-                                                       self.conservative_density,
                                                        ibp=ibp)
+
+                # Store the conservative transport form in the mass_weighted label,
+                # but by default use an advective form.
+                form = mass_weighted(advective_form, mass_weighted_form)
             else:
                 raise NotImplementedError('Upwind transport scheme has not been '
                                           + 'implemented for this transport equation type')
-
         self.form = form
 
 

gusto/time_discretisation/time_discretisation.py

@@ -17,7 +17,7 @@
 from firedrake.utils import cached_property
 
 from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions
-from gusto.core.labels import time_derivative, prognostic, physics_label
+from gusto.core.labels import time_derivative, prognostic, physics_label, mass_weighted
 from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.time_discretisation.wrappers import *
 
@@ -161,6 +161,28 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     self.evaluate_source.append(t.labels[physics_name])
                     self.physics_names.append(t.labels[physics_name])
 
+        # Check if there are any mass-weighted terms:
+        if len(self.residual.label_map(lambda t: t.has_label(mass_weighted), map_if_false=drop)) > 0:
+            for field in equation.field_names:
+
+                # Check if the mass term for this prognostic is mass-weighted
+                if len(self.residual.label_map(lambda t: t.get(prognostic) == field and t.has_label(time_derivative) and t.has_label(mass_weighted), map_if_false=drop)) == 1:
+                    field_terms = self.residual.label_map(lambda t: t.get(prognostic) == field and not t.has_label(time_derivative), map_if_false=drop)
+
+                    # Check that the equation for this prognostic does not involve
+                    # both mass-weighted and non-mass-weighted terms; if so, a split
+                    # timestepper should be used instead.
+                    if len(field_terms.label_map(lambda t: t.has_label(mass_weighted), map_if_false=drop)) > 0:
+                        if len(field_terms.label_map(lambda t: not t.has_label(mass_weighted), map_if_false=drop)) > 0:
+                            raise ValueError(f"Mass-weighted and non-mass-weighted terms are present in a timestepping equation for {field}. As these terms cannot be solved for simultaneously, a split timestepping method should be used instead.")
+                        else:
+                            # Replace the terms with a mass_weighted label with the
+                            # mass_weighted form. It is important that the labels from
+                            # this new form are used.
+                            self.residual = self.residual.label_map(
+                                lambda t: t.get(prognostic) == field and t.has_label(mass_weighted),
+                                map_if_true=lambda t: t.get(mass_weighted))
+
         # -------------------------------------------------------------------- #
         # Set up Wrappers
         # -------------------------------------------------------------------- #

integration-tests/model/test_conservative_transport_with_physics.py

@@ -0,0 +1,97 @@
+"""
+This tests that conservative transport is correctly working with a physics
+scheme. The conservative transport equations require the tracer to be multiplied
+by the density through the 'mass_weighted' label, whilst the physics terms
+do not. Here, we test that we correctly replace the transport terms with the
+mass_weighted counterpart but leave the physics terms unchanged.
+"""
+
+from gusto import *
+from firedrake import (as_vector, PeriodicSquareMesh, SpatialCoordinate,
+                       assemble, Constant, conditional, sin, pi)
+
+
+def run_conservative_transport_with_physics(dirname):
+
+    # ------------------------------------------------------------------------ #
+    # Set up model objects
+    # ------------------------------------------------------------------------ #
+
+    # set up mesh and domain
+    L = 10
+    nx = 10
+    mesh = PeriodicSquareMesh(nx, nx, L, quadrilateral=True)
+    dt = 0.1
+    tmax = 5*dt
+    domain = Domain(mesh, dt, "RTCF", 1)
+    x, y = SpatialCoordinate(mesh)
+
+    rho_d_space = 'DG'
+    ash_space = 'DG'
+
+    ash = ActiveTracer(name='ash', space=ash_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)
+
+    tracers = [ash, rho_d]
+
+    eqn = CoupledTransportEquation(domain, active_tracers=tracers)
+
+    # I/O
+    output = OutputParameters(dirname=dirname+"/conservative_transport_with_physics",
+                              dumpfreq=1)
+    diagnostic_fields = [CourantNumber()]
+    io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+    transport_method = [DGUpwind(eqn, "rho_d"), DGUpwind(eqn, "ash")]
+
+    # Physics scheme --------------------------------------------------------- #
+    # Source is a constant, but constrained to a box in the bottom left corner
+    # of size 1-by-1, such that the total ash value
+    # should be equal to tmax = 0.5.
+    basic_expression = conditional(x < 1.0, conditional(y < 1.0, -Constant(1.0), Constant(0.0)), Constant(0.0))
+
+    physics_schemes = [(SourceSink(eqn, 'ash', basic_expression), SSPRK3(domain))]
+
+    # Time stepper
+    stepper = SplitPrescribedTransport(eqn, SSPRK3(domain, increment_form=False),
+                                       io, transport_method,
+                                       physics_schemes=physics_schemes)
+
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+
+    rho0 = stepper.fields("rho_d")
+    ash0 = stepper.fields("ash")
+
+    # Set a spatially varying density field and no ash
+    rho0.interpolate(1000.0*sin(pi*x/L)*sin(pi*y/L)+1000.0)
+    ash0.interpolate(Constant(0.0))
+
+    # Constant wind
+    u = stepper.fields("u")
+    u.project(as_vector([0.5, 0.5]))
+
+    # ------------------------------------------------------------------------ #
+    # Run
+    # ------------------------------------------------------------------------ #
+
+    stepper.run(t=0, tmax=tmax)
+    return stepper
+
+
+def test_conservative_transport_with_physics(tmpdir):
+    dirname = str(tmpdir)
+    stepper = run_conservative_transport_with_physics(dirname)
+    final_ash = stepper.fields("ash")
+
+    final_total_ash = assemble(final_ash*dx)
+
+    tol = 1e-3
+    assert np.abs(final_total_ash - 0.5) < tol, \
+        "Conservative transport did not correctly implement the Source physics"