Merge pull request #537 from firedrakeproject/split_timestepper

Split timestepper

gusto/equations/shallow_water_equations.py

@@ -189,14 +189,15 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
                                           self.X)
             residual += topography_form
 
-        # thermal source terms not involving topography
+        # thermal source terms not involving topography.
+        # label these as the equivalent pressure gradient term
         if self.thermal:
             n = FacetNormal(domain.mesh)
-            source_form = subject(prognostic(-D*div(b*w)*dx
-                                             - 0.5*b*div(D*w)*dx
-                                             + jump(b*w, n)*avg(D)*dS
-                                             + 0.5*jump(D*w, n)*avg(b)*dS,
-                                             'u'), self.X)
+            source_form = pressure_gradient(subject(prognostic(-D*div(b*w)*dx
+                                                    - 0.5*b*div(D*w)*dx
+                                                    + jump(b*w, n)*avg(D)*dS
+                                                    + 0.5*jump(D*w, n)*avg(b)*dS,
+                                                    'u'), self.X))
             residual += source_form
 
         # -------------------------------------------------------------------- #

gusto/timestepping/split_timestepper.py

@@ -1,13 +1,165 @@
 """Split timestepping methods for generically solving terms separately."""
 
 from firedrake import Projector
-from firedrake.fml import Label
+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 Timestepper
+from gusto.timestepping.timestepper import BaseTimestepper, Timestepper
+from numpy import ones
 
-__all__ = ["SplitPhysicsTimestepper", "SplitPrescribedTransport"]
+__all__ = ["SplitTimestepper", "SplitPhysicsTimestepper", "SplitPrescribedTransport"]
+
+
+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')
+
+    def setup_fields(self):
+        self.x = TimeLevelFields(self.equation, 1)
+        self.fields = StateFields(self.x, self.equation.prescribed_fields,
+                                  *self.io.output.dumplist)
+
+    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)
+
+    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))
 
 
 class SplitPhysicsTimestepper(Timestepper):

integration-tests/equations/test_boussinesq.py

@@ -128,6 +128,8 @@ def test_boussinesq(tmpdir, compressible):
         diff_array = new_variable.dat.data - check_variable.dat.data
         error = np.linalg.norm(diff_array) / np.linalg.norm(check_variable.dat.data)
 
+        test_type = 'compressible' if compressible else 'incompressible'
+
         # Slack values chosen to be robust to different platforms
         assert error < 1e-10, f'Values for {variable} in ' + \
-            'Incompressible test do not match KGO values'
+            f'{test_type} test do not match KGO values'

integration-tests/model/test_split_timestepper.py

@@ -0,0 +1,127 @@
+"""
+This script tests the split_timestepper using an advection-diffusion
+equation with a physics parametrisation. Three different splittings are
+tested, including splitting the dynamics and physics into two substeps
+with different timestep sizes.
+"""
+
+from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, as_vector,
+                       norm, Constant, conditional, sqrt, VectorFunctionSpace)
+from gusto import *
+import pytest
+
+
+def run_split_timestepper_adv_diff_physics(tmpdir, timestepper):
+
+    # ------------------------------------------------------------------------ #
+    # Set up model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    dt = 0.02
+    tmax = 1.0
+    L = 10
+    mesh = PeriodicIntervalMesh(20, L)
+    domain = Domain(mesh, dt, "CG", 1)
+
+    # Equation
+    diffusion_params = DiffusionParameters(kappa=0.75, mu=5)
+    V = domain.spaces("DG")
+    Vu = VectorFunctionSpace(mesh, "CG", 1)
+
+    equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu,
+                                          diffusion_parameters=diffusion_params)
+    spatial_methods = [DGUpwind(equation, "f"),
+                       InteriorPenaltyDiffusion(equation, "f", diffusion_params)]
+
+    x = SpatialCoordinate(mesh)
+
+    # Add a source term to inject mass into the domain.
+    # Without the diffusion, this would simply add 0.1
+    # units of mass equally across the domain.
+    source_expression = -Constant(0.1)
+
+    physics_schemes = [(SourceSink(equation, "f", source_expression), SSPRK3(domain))]
+
+    # I/O
+    output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
+    io = IO(domain, output)
+
+    # Time stepper
+    if timestepper == 'split1':
+        # Split with no defined weights
+        dynamics_schemes = {'transport': ImplicitMidpoint(domain),
+                            'diffusion': ForwardEuler(domain)}
+        term_splitting = ['transport', 'diffusion', 'physics']
+        stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
+                                   io, spatial_methods=spatial_methods,
+                                   physics_schemes=physics_schemes)
+    elif timestepper == 'split2':
+        # Transport split into two substeps
+        dynamics_schemes = {'transport': SSPRK3(domain),
+                            'diffusion': ForwardEuler(domain)}
+        term_splitting = ['diffusion', 'transport', 'physics', 'transport']
+        weights = [1., 0.6, 1., 0.4]
+        stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
+                                   io, weights=weights, spatial_methods=spatial_methods,
+                                   physics_schemes=physics_schemes)
+    else:
+        # Physics split into two substeps
+        dynamics_schemes = {'transport': SSPRK3(domain),
+                            'diffusion': SSPRK3(domain)}
+        term_splitting = ['physics', 'transport', 'diffusion', 'physics']
+        weights = [1./3., 1., 1., 2./3.]
+        stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
+                                   io, weights=weights, spatial_methods=spatial_methods,
+                                   physics_schemes=physics_schemes)
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+
+    xc_init = 0.25*L
+    xc_end = 0.75*L
+    umax = 0.5*L/tmax
+
+    # Get minimum distance on periodic interval to xc
+    x_init = conditional(sqrt((x[0] - xc_init)**2) < 0.5*L,
+                         x[0] - xc_init, L + x[0] - xc_init)
+
+    x_end = conditional(sqrt((x[0] - xc_end)**2) < 0.5*L,
+                        x[0] - xc_end, L + x[0] - xc_end)
+
+    f_init = 5.0
+    f_end = f_init / 2.0
+    f_width_init = L / 10.0
+    f_width_end = f_width_init * 2.0
+    f_init_expr = f_init*exp(-(x_init / f_width_init)**2)
+
+    # The end Gaussian should be advected by half the domain
+    # length, be more spread out due to the dissipation,
+    # and includes more mass due to the source term.
+    f_end_expr = 0.1 + f_end*exp(-(x_end / f_width_end)**2)
+
+    stepper.fields('f').interpolate(f_init_expr)
+    stepper.fields('u').interpolate(as_vector([Constant(umax)]))
+    f_end = stepper.fields('f_end', space=V)
+    f_end.interpolate(f_end_expr)
+
+    # ------------------------------------------------------------------------ #
+    # Run
+    # ------------------------------------------------------------------------ #
+
+    stepper.run(0, tmax=tmax)
+
+    error = norm(stepper.fields('f') - f_end) / norm(f_end)
+
+    return error
+
+
+@pytest.mark.parametrize("timestepper", ["split1", "split2", "split3"])
+def test_split_timestepper_adv_diff_physics(tmpdir, timestepper):
+
+    tol = 0.015
+    error = run_split_timestepper_adv_diff_physics(tmpdir, timestepper)
+    print(error)
+    assert error < tol, 'The split timestepper in the advection-diffusion' + \
+        'equation with source physics has an error greater than ' + \
+        'the permitted tolerance'