diff --git a/gusto/physics.py b/gusto/physics.py index 7bebece17..934d508b4 100644 --- a/gusto/physics.py +++ b/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 diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 7c1c2a441..0ebefd8d1 100644 --- a/gusto/timeloop.py +++ b/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,