PR #457: from firedrakeproject/terminator_toy

Add terminator toy physics parametrisation, and split prescribed transport timestepper

gusto/physics.py

@@ -32,7 +32,7 @@
 __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain",
            "AdvectedMoments", "InstantRain", "SWSaturationAdjustment",
            "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment",
-           "SuppressVerticalWind", "BoundaryLayerMixing"]
+           "SuppressVerticalWind", "BoundaryLayerMixing", "TerminatorToy"]
 
 
 class PhysicsParametrisation(object, metaclass=ABCMeta):
@@ -1751,3 +1751,76 @@ def evaluate(self, x_in, dt):
 
         self.X.assign(x_in)
         self.rho_recoverer.project()
+
+
+class TerminatorToy(PhysicsParametrisation):
+    """
+    Setup the Terminator Toy chemistry interaction
+    as specified in 'The terminator toy chemistry test ...'
+    Lauritzen et. al. (2014).
+
+    The coupled equations for the two species are given by:
+
+    D/Dt (X) = 2Kx
+    D/Dt (X2) = -Kx
+
+    where Kx = k1*X2 - k2*(X**2)
+    """
+
+    def __init__(self, equation, k1=1, k2=1,
+                 species1_name='X', species2_name='X2'):
+        """
+        Args:
+            equation (:class: 'PrognosticEquationSet'): the model's equation
+            k1(float, optional): Reaction rate for species 1 (X). Defaults to a constant
+              over the domain.
+            k2(float, optional): Reaction rate for species 2 (X2). Defaults to a constant
+              over the domain.
+            species1_name(str, optional): Name of the first interacting species. Defaults
+              to 'X'.
+            species2_name(str, optional): Name of the second interacting species. Defaults
+              to 'X2'.
+        """
+
+        label_name = 'terminator_toy'
+        super().__init__(equation, label_name, parameters=None)
+
+        if species1_name not in equation.field_names:
+            raise ValueError(f"Field {species1_name} does not exist in the equation set")
+        if species2_name not in equation.field_names:
+            raise ValueError(f"Field {species2_name} does not exist in the equation set")
+
+        self.species1_idx = equation.field_names.index(species1_name)
+        self.species2_idx = equation.field_names.index(species2_name)
+
+        assert equation.function_space.sub(self.species1_idx) == equation.function_space.sub(self.species2_idx), "The function spaces for the two species need to be the same"
+
+        self.Xq = Function(equation.X.function_space())
+        Xq = self.Xq
+
+        species1 = split(Xq)[self.species1_idx]
+        species2 = split(Xq)[self.species2_idx]
+
+        test_1 = equation.tests[self.species1_idx]
+        test_2 = equation.tests[self.species2_idx]
+
+        Kx = k1*species2 - k2*(species1**2)
+
+        source1_expr = test_1 * 2*Kx * dx
+        source2_expr = test_2 * -Kx * dx
+
+        equation.residual -= self.label(subject(prognostic(source1_expr, 'X'), Xq), self.evaluate)
+        equation.residual -= self.label(subject(prognostic(source2_expr, 'X2'), Xq), self.evaluate)
+
+    def evaluate(self, x_in, dt):
+        """
+        Evaluates the source/sink for the coalescence process.
+
+        Args:
+            x_in (:class:`Function`): the (mixed) field to be evolved.
+            dt (:class:`Constant`): the time interval for the scheme.
+        """
+
+        logger.info(f'Evaluating physics parametrisation {self.label.label}')
+
+        pass

gusto/timeloop.py

@@ -17,8 +17,8 @@
 from gusto.transport_methods import TransportMethod
 import ufl
 
-__all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton",
-           "PrescribedTransport"]
+__all__ = ["Timestepper", "SplitPhysicsTimestepper", "SplitPrescribedTransport",
+           "SemiImplicitQuasiNewton", "PrescribedTransport"]
 
 
 class BaseTimestepper(object, metaclass=ABCMeta):
@@ -124,7 +124,6 @@ def setup_transporting_velocity(self, scheme):
                 transport term should be replaced with the transport term of
                 this discretisation.
         """
-
         if self.transporting_velocity == "prognostic" and "u" in self.fields._field_names:
             # Use the prognostic wind variable as the transporting velocity
             u_idx = self.equation.field_names.index('u')
@@ -378,6 +377,91 @@ def timestep(self):
                 scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name))
 
 
+class SplitPrescribedTransport(Timestepper):
+    """
+    Implements a timeloop where the physics terms are solved separately from
+    the dynamics, like with SplitPhysicsTimestepper, but here we define
+    a prescribed transporting velocity, as opposed to having the
+    velocity as a prognostic variable.
+    """
+
+    def __init__(self, equation, scheme, io, spatial_methods=None,
+                 physics_schemes=None,
+                 prescribed_transporting_velocity=None):
+        """
+        Args:
+            equation (:class:`PrognosticEquation`): the prognostic equation
+            scheme (:class:`TimeDiscretisation`): the scheme to use to timestep
+                the prognostic equation
+            io (:class:`IO`): the model's object for controlling input/output.
+            spatial_methods (iter,optional): a list of objects describing the
+                methods to use for discretising transport or diffusion terms
+                for each transported/diffused variable. Defaults to None,
+                in which case the terms follow the original discretisation in
+                the equation.
+            physics_schemes: (list, optional): a list of tuples of the form
+                (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`),
+                pairing physics parametrisations and timestepping schemes to use
+                for each. Timestepping schemes for physics can be explicit
+                or implicit. Defaults to None.
+            prescribed_transporting_velocity: (field, optional): A known
+                velocity field that is used for the transport of tracers.
+                This can be made time-varying by defining a python function
+                that uses time as an argument.
+                Defaults to None.
+        """
+
+        # As we handle physics differently to the Timestepper, these are not
+        # passed to the super __init__
+        super().__init__(equation, scheme, io, spatial_methods=spatial_methods)
+
+        if physics_schemes is not None:
+            self.physics_schemes = physics_schemes
+        else:
+            self.physics_schemes = []
+
+        for parametrisation, phys_scheme in self.physics_schemes:
+            # check that the supplied schemes for physics are valid
+            if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only:
+                assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \
+                    ("Only explicit time discretisations can be used with "
+                     + f"physics scheme {parametrisation.label.label}")
+            apply_bcs = False
+            phys_scheme.setup(equation, apply_bcs, parametrisation.label)
+
+        if prescribed_transporting_velocity is not None:
+            self.velocity_projection = Projector(
+                prescribed_transporting_velocity(self.t),
+                self.fields('u'))
+        else:
+            self.velocity_projection = None
+
+    @property
+    def transporting_velocity(self):
+        return self.fields('u')
+
+    def setup_scheme(self):
+        self.setup_equation(self.equation)
+        # 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)
+        apply_bcs = True
+        self.scheme.setup(self.equation, apply_bcs, dynamics)
+        self.setup_transporting_velocity(self.scheme)
+        self.scheme.courant_max = self.io.courant_max
+
+    def timestep(self):
+
+        if self.velocity_projection is not None:
+            self.velocity_projection.project()
+
+        super().timestep()
+
+        with timed_stage("Physics"):
+            for _, scheme in self.physics_schemes:
+                scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name))
+
+
 class SemiImplicitQuasiNewton(BaseTimestepper):
     """
     Implements a semi-implicit quasi-Newton discretisation,