"""Split timestepping methods for generically solving terms separately."""
from firedrake import Projector
from firedrake.fml import Label, drop
from pyop2.profiling import timed_stage
from gusto.core import TimeLevelFields, StateFields
from gusto.core.labels import time_derivative, physics_label
from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation
from gusto.timestepping.timestepper import BaseTimestepper, Timestepper
from numpy import ones
__all__ = ["SplitTimestepper", "SplitPhysicsTimestepper", "SplitPrescribedTransport"]
[docs]
class SplitTimestepper(BaseTimestepper):
"""
Implements a timeloop by applying separate schemes to different terms, e.g, physics
and individual dynamics components in a user-defined order. This allows a different
time discretisation to be applied to each defined component. Different terms can be
substepped by specifying weights; this enables Strang-Splitting to be applied.
"""
def __init__(self, equation, term_splitting, dynamics_schemes, io,
weights=None, spatial_methods=None, physics_schemes=None):
"""
Args:
equation (:class:`PrognosticEquation`): the prognostic equation
term_splitting (list): a list of labels specifying the terms that should
be solved separately and the order to do so.
dynamics_schemes: (:class:`TimeDiscretisation`) A list of time
discretisations for the dynamics schemes. A scheme must be
provided for each non-physics label that is provided in the
term_splitting list.
io (:class:`IO`): the model's object for controlling input/output.
weights (array, optional): An array of weights for substepping
of any dynamics or physics scheme. The sum of weights for
each distinct label in term_splitting must be 1.
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 must be explicit.
Defaults to None.
"""
if spatial_methods is not None:
self.spatial_methods = spatial_methods
else:
self.spatial_methods = []
# If we have physics schemes, extract these.
if 'physics' in term_splitting:
if physics_schemes is None:
raise ValueError('Physics schemes need to be specified when applying '
+ 'a physics splitting in the SplitTimestepper')
else:
# Check that the weights are correct for physics:
count = 0
if weights is not None:
for idx, term in enumerate(term_splitting):
if term == 'physics':
count += weights[idx]
if count != 1:
raise ValueError('Incorrect weights are specified for the '
+ 'physics schemes in the split timestepper.')
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}")
self.term_splitting = term_splitting
self.dynamics_schemes = dynamics_schemes
if weights is not None:
self.weights = weights
else:
self.weights = ones(len(self.term_splitting))
# Check that each dynamics label in term_splitting has a corresponding
# dynamics scheme
for term in self.term_splitting:
if term != 'physics':
assert term in self.dynamics_schemes, \
f'The {term} terms do not have a specified scheme in the split timestepper'
# Multilevel schemes are currently not supported for the dynamics terms.
for label, scheme in self.dynamics_schemes.items():
assert scheme.nlevels == 1, \
"Multilevel schemes are not currently implemented in the split timestepper"
# As we handle physics in separate parametrisations, these are not
# passed to the super __init__
super().__init__(equation, io)
# Check that each dynamics term is specified by a label
# in the term_splitting list, but also that there are not
# multiple labels, i.e. there is a single specified time discretisation.
# When using weights, these should add to 1 for each term.
terms = self.equation.residual.label_map(
lambda t: any(t.has_label(time_derivative, physics_label)), map_if_true=drop
)
for term in terms:
count = 0
for idx, label in enumerate(self.term_splitting):
if term.has_label(Label(label)):
count += self.weights[idx]
if count != 1:
raise ValueError('The term_splitting list does not correctly cover '
+ 'the dynamics terms in the equation(s).')
# Timesteps for each scheme in the term_splitting list
self.split_dts = [self.equation.domain.dt*weight for weight in self.weights]
@property
def transporting_velocity(self):
return self.fields('u')
[docs]
def setup_fields(self):
self.x = TimeLevelFields(self.equation, 1)
self.fields = StateFields(self.x, self.equation.prescribed_fields,
*self.io.output.dumplist)
[docs]
def setup_scheme(self):
"""Sets up transport, diffusion and physics schemes"""
# TODO: apply_bcs should be False for advection but this means
# tests with KGOs fail
self.setup_equation(self.equation)
apply_bcs = True
for label, scheme in self.dynamics_schemes.items():
scheme.setup(self.equation, apply_bcs, Label(label))
self.setup_transporting_velocity(scheme)
if self.io.output.log_courant and label == 'transport':
scheme.courant_max = self.io.courant_max
apply_bcs = False
for parametrisation, scheme in self.physics_schemes:
scheme.setup(self.equation, apply_bcs, parametrisation.label)
[docs]
def timestep(self):
for idx, term in enumerate(self.term_splitting):
split_dt = self.split_dts[idx]
if term == 'physics':
with timed_stage("Physics"):
for _, scheme in self.physics_schemes:
scheme.dt = split_dt
scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name))
else:
scheme = self.dynamics_schemes[term]
scheme.dt = split_dt
scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name))
[docs]
class SplitPhysicsTimestepper(Timestepper):
"""
Implements a timeloop by applying schemes separately to the physics and
dynamics. This 'splits' the physics from the dynamics and allows a different
scheme to be applied to the physics terms than the prognostic equation.
"""
def __init__(self, equation, scheme, io, spatial_methods=None,
physics_schemes=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 must be explicit.
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)
@property
def transporting_velocity(self):
return "prognostic"
[docs]
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)
if self.io.output.log_courant:
self.scheme.courant_max = self.io.courant_max
[docs]
def timestep(self):
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))
[docs]
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, prescribed_transporting_velocity,
spatial_methods=None, physics_schemes=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.
prescribed_transporting_velocity: (bool): Whether a time-varying
transporting velocity will be defined. If True, this will
require the transporting velocity to be setup by calling either
the `setup_prescribed_expr` or `setup_prescribed_apply` methods.
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)
self.prescribed_transport_velocity = prescribed_transporting_velocity
self.is_velocity_setup = not self.prescribed_transport_velocity
self.velocity_projection = None
self.velocity_apply = None
@property
def transporting_velocity(self):
return self.fields('u')
[docs]
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)
if self.io.output.log_courant:
self.scheme.courant_max = self.io.courant_max
[docs]
def setup_prescribed_expr(self, expr_func):
"""
Sets up the prescribed transporting velocity, through a python function
which has time as an argument, and returns a `ufl.Expr`. This allows the
velocity to be updated with time.
Args:
expr_func (func): a python function with a single argument that
represents the model time, and returns a `ufl.Expr`.
"""
if self.is_velocity_setup:
raise RuntimeError('Prescribed velocity already set up!')
self.velocity_projection = Projector(
expr_func(self.t), self.fields('u')
)
self.is_velocity_setup = True
[docs]
def setup_prescribed_apply(self, apply_method):
"""
Sets up the prescribed transporting velocity, through a python function
which has time as an argument. This function will perform the evaluation
of the transporting velocity.
Args:
expr_func (func): a python function with a single argument that
represents the model time, and performs the evaluation of the
transporting velocity.
"""
if self.is_velocity_setup:
raise RuntimeError('Prescribed velocity already set up!')
self.velocity_apply = apply_method
self.is_velocity_setup = True
[docs]
def run(self, t, tmax, pick_up=False):
"""
Runs the model for the specified time, from t to tmax
Args:
t (float): the start time of the run
tmax (float): the end time of the run
pick_up: (bool): specify whether to pick_up from a previous run
"""
# Throw an error if no transporting velocity has been set up
if self.prescribed_transport_velocity and not self.is_velocity_setup:
raise RuntimeError(
'A time-varying prescribed velocity is required. This must be '
+ 'set up through calling either the setup_prescribed_expr or '
+ 'setup_prescribed_apply routines.')
# It's best to have evaluated the velocity before we start
if self.velocity_projection is not None:
self.velocity_projection.project()
if self.velocity_apply is not None:
self.velocity_apply(self.t)
super().run(t, tmax, pick_up=pick_up)
[docs]
def timestep(self):
if self.velocity_projection is not None:
self.velocity_projection.project()
if self.velocity_apply is not None:
self.velocity_apply(self.t)
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))