Skip to content

Commit 3a5367e

Browse files
tommbendalltommbendall
authored
Merge pull request #537 from firedrakeproject/split_timestepper
Split timestepper
2 parents 471e228 + fb8ef9c commit 3a5367e

File tree

4 files changed

+292
-10
lines changed

4 files changed

+292
-10
lines changed

‎gusto/equations/shallow_water_equations.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -189,14 +189,15 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
189189
self.X)
190190
residual += topography_form
191191

192-
# thermal source terms not involving topography
192+
# thermal source terms not involving topography.
193+
# label these as the equivalent pressure gradient term
193194
if self.thermal:
194195
n = FacetNormal(domain.mesh)
195-
source_form = subject(prognostic(-D*div(b*w)*dx
196-
- 0.5*b*div(D*w)*dx
197-
+ jump(b*w, n)*avg(D)*dS
198-
+ 0.5*jump(D*w, n)*avg(b)*dS,
199-
'u'), self.X)
196+
source_form = pressure_gradient(subject(prognostic(-D*div(b*w)*dx
197+
- 0.5*b*div(D*w)*dx
198+
+ jump(b*w, n)*avg(D)*dS
199+
+ 0.5*jump(D*w, n)*avg(b)*dS,
200+
'u'), self.X))
200201
residual += source_form
201202

202203
# -------------------------------------------------------------------- #

‎gusto/timestepping/split_timestepper.py

Lines changed: 155 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,165 @@
11
"""Split timestepping methods for generically solving terms separately."""
22

33
from firedrake import Projector
4-
from firedrake.fml import Label
4+
from firedrake.fml import Label, drop
55
from pyop2.profiling import timed_stage
6+
from gusto.core import TimeLevelFields, StateFields
67
from gusto.core.labels import time_derivative, physics_label
78
from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation
8-
from gusto.timestepping.timestepper import Timestepper
9+
from gusto.timestepping.timestepper import BaseTimestepper, Timestepper
10+
from numpy import ones
911

10-
__all__ = ["SplitPhysicsTimestepper", "SplitPrescribedTransport"]
12+
__all__ = ["SplitTimestepper", "SplitPhysicsTimestepper", "SplitPrescribedTransport"]
13+
14+
15+
class SplitTimestepper(BaseTimestepper):
16+
"""
17+
Implements a timeloop by applying separate schemes to different terms, e.g, physics
18+
and individual dynamics components in a user-defined order. This allows a different
19+
time discretisation to be applied to each defined component. Different terms can be
20+
substepped by specifying weights; this enables Strang-Splitting to be applied.
21+
"""
22+
23+
def __init__(self, equation, term_splitting, dynamics_schemes, io,
24+
weights=None, spatial_methods=None, physics_schemes=None):
25+
"""
26+
Args:
27+
equation (:class:`PrognosticEquation`): the prognostic equation
28+
term_splitting (list): a list of labels specifying the terms that should
29+
be solved separately and the order to do so.
30+
dynamics_schemes: (:class:`TimeDiscretisation`) A list of time
31+
discretisations for the dynamics schemes. A scheme must be
32+
provided for each non-physics label that is provided in the
33+
term_splitting list.
34+
io (:class:`IO`): the model's object for controlling input/output.
35+
weights (array, optional): An array of weights for substepping
36+
of any dynamics or physics scheme. The sum of weights for
37+
each distinct label in term_splitting must be 1.
38+
spatial_methods (iter,optional): a list of objects describing the
39+
methods to use for discretising transport or diffusion terms
40+
for each transported/diffused variable. Defaults to None,
41+
in which case the terms follow the original discretisation in
42+
the equation.
43+
physics_schemes: (list, optional): a list of tuples of the form
44+
(:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`),
45+
pairing physics parametrisations and timestepping schemes to use
46+
for each. Timestepping schemes for physics must be explicit.
47+
Defaults to None.
48+
"""
49+
50+
if spatial_methods is not None:
51+
self.spatial_methods = spatial_methods
52+
else:
53+
self.spatial_methods = []
54+
55+
# If we have physics schemes, extract these.
56+
if 'physics' in term_splitting:
57+
if physics_schemes is None:
58+
raise ValueError('Physics schemes need to be specified when applying '
59+
+ 'a physics splitting in the SplitTimestepper')
60+
else:
61+
# Check that the weights are correct for physics:
62+
count = 0
63+
if weights is not None:
64+
for idx, term in enumerate(term_splitting):
65+
if term == 'physics':
66+
count += weights[idx]
67+
if count != 1:
68+
raise ValueError('Incorrect weights are specified for the '
69+
+ 'physics schemes in the split timestepper.')
70+
self.physics_schemes = physics_schemes
71+
else:
72+
self.physics_schemes = []
73+
74+
for parametrisation, phys_scheme in self.physics_schemes:
75+
# Check that the supplied schemes for physics are valid
76+
if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only:
77+
assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \
78+
("Only explicit time discretisations can be used with "
79+
+ f"physics scheme {parametrisation.label.label}")
80+
81+
self.term_splitting = term_splitting
82+
self.dynamics_schemes = dynamics_schemes
83+
84+
if weights is not None:
85+
self.weights = weights
86+
else:
87+
self.weights = ones(len(self.term_splitting))
88+
89+
# Check that each dynamics label in term_splitting has a corresponding
90+
# dynamics scheme
91+
for term in self.term_splitting:
92+
if term != 'physics':
93+
assert term in self.dynamics_schemes, \
94+
f'The {term} terms do not have a specified scheme in the split timestepper'
95+
96+
# Multilevel schemes are currently not supported for the dynamics terms.
97+
for label, scheme in self.dynamics_schemes.items():
98+
assert scheme.nlevels == 1, \
99+
"Multilevel schemes are not currently implemented in the split timestepper"
100+
101+
# As we handle physics in separate parametrisations, these are not
102+
# passed to the super __init__
103+
super().__init__(equation, io)
104+
105+
# Check that each dynamics term is specified by a label
106+
# in the term_splitting list, but also that there are not
107+
# multiple labels, i.e. there is a single specified time discretisation.
108+
# When using weights, these should add to 1 for each term.
109+
terms = self.equation.residual.label_map(
110+
lambda t: any(t.has_label(time_derivative, physics_label)), map_if_true=drop
111+
)
112+
for term in terms:
113+
count = 0
114+
for idx, label in enumerate(self.term_splitting):
115+
if term.has_label(Label(label)):
116+
count += self.weights[idx]
117+
if count != 1:
118+
raise ValueError('The term_splitting list does not correctly cover '
119+
+ 'the dynamics terms in the equation(s).')
120+
121+
# Timesteps for each scheme in the term_splitting list
122+
self.split_dts = [self.equation.domain.dt*weight for weight in self.weights]
123+
124+
@property
125+
def transporting_velocity(self):
126+
return self.fields('u')
127+
128+
def setup_fields(self):
129+
self.x = TimeLevelFields(self.equation, 1)
130+
self.fields = StateFields(self.x, self.equation.prescribed_fields,
131+
*self.io.output.dumplist)
132+
133+
def setup_scheme(self):
134+
"""Sets up transport, diffusion and physics schemes"""
135+
# TODO: apply_bcs should be False for advection but this means
136+
# tests with KGOs fail
137+
self.setup_equation(self.equation)
138+
139+
apply_bcs = True
140+
for label, scheme in self.dynamics_schemes.items():
141+
scheme.setup(self.equation, apply_bcs, Label(label))
142+
self.setup_transporting_velocity(scheme)
143+
if self.io.output.log_courant and label == 'transport':
144+
scheme.courant_max = self.io.courant_max
145+
146+
apply_bcs = False
147+
for parametrisation, scheme in self.physics_schemes:
148+
scheme.setup(self.equation, apply_bcs, parametrisation.label)
149+
150+
def timestep(self):
151+
152+
for idx, term in enumerate(self.term_splitting):
153+
split_dt = self.split_dts[idx]
154+
if term == 'physics':
155+
with timed_stage("Physics"):
156+
for _, scheme in self.physics_schemes:
157+
scheme.dt = split_dt
158+
scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name))
159+
else:
160+
scheme = self.dynamics_schemes[term]
161+
scheme.dt = split_dt
162+
scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name))
11163

12164

13165
class SplitPhysicsTimestepper(Timestepper):

‎integration-tests/equations/test_boussinesq.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,8 @@ def test_boussinesq(tmpdir, compressible):
128128
diff_array = new_variable.dat.data - check_variable.dat.data
129129
error = np.linalg.norm(diff_array) / np.linalg.norm(check_variable.dat.data)
130130

131+
test_type = 'compressible' if compressible else 'incompressible'
132+
131133
# Slack values chosen to be robust to different platforms
132134
assert error < 1e-10, f'Values for {variable} in ' + \
133-
'Incompressible test do not match KGO values'
135+
f'{test_type} test do not match KGO values'

‎integration-tests/model/test_split_timestepper.py

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
"""
2+
This script tests the split_timestepper using an advection-diffusion
3+
equation with a physics parametrisation. Three different splittings are
4+
tested, including splitting the dynamics and physics into two substeps
5+
with different timestep sizes.
6+
"""
7+
8+
from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, as_vector,
9+
norm, Constant, conditional, sqrt, VectorFunctionSpace)
10+
from gusto import *
11+
import pytest
12+
13+
14+
def run_split_timestepper_adv_diff_physics(tmpdir, timestepper):
15+
16+
# ------------------------------------------------------------------------ #
17+
# Set up model objects
18+
# ------------------------------------------------------------------------ #
19+
20+
# Domain
21+
dt = 0.02
22+
tmax = 1.0
23+
L = 10
24+
mesh = PeriodicIntervalMesh(20, L)
25+
domain = Domain(mesh, dt, "CG", 1)
26+
27+
# Equation
28+
diffusion_params = DiffusionParameters(kappa=0.75, mu=5)
29+
V = domain.spaces("DG")
30+
Vu = VectorFunctionSpace(mesh, "CG", 1)
31+
32+
equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu,
33+
diffusion_parameters=diffusion_params)
34+
spatial_methods = [DGUpwind(equation, "f"),
35+
InteriorPenaltyDiffusion(equation, "f", diffusion_params)]
36+
37+
x = SpatialCoordinate(mesh)
38+
39+
# Add a source term to inject mass into the domain.
40+
# Without the diffusion, this would simply add 0.1
41+
# units of mass equally across the domain.
42+
source_expression = -Constant(0.1)
43+
44+
physics_schemes = [(SourceSink(equation, "f", source_expression), SSPRK3(domain))]
45+
46+
# I/O
47+
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
48+
io = IO(domain, output)
49+
50+
# Time stepper
51+
if timestepper == 'split1':
52+
# Split with no defined weights
53+
dynamics_schemes = {'transport': ImplicitMidpoint(domain),
54+
'diffusion': ForwardEuler(domain)}
55+
term_splitting = ['transport', 'diffusion', 'physics']
56+
stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
57+
io, spatial_methods=spatial_methods,
58+
physics_schemes=physics_schemes)
59+
elif timestepper == 'split2':
60+
# Transport split into two substeps
61+
dynamics_schemes = {'transport': SSPRK3(domain),
62+
'diffusion': ForwardEuler(domain)}
63+
term_splitting = ['diffusion', 'transport', 'physics', 'transport']
64+
weights = [1., 0.6, 1., 0.4]
65+
stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
66+
io, weights=weights, spatial_methods=spatial_methods,
67+
physics_schemes=physics_schemes)
68+
else:
69+
# Physics split into two substeps
70+
dynamics_schemes = {'transport': SSPRK3(domain),
71+
'diffusion': SSPRK3(domain)}
72+
term_splitting = ['physics', 'transport', 'diffusion', 'physics']
73+
weights = [1./3., 1., 1., 2./3.]
74+
stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
75+
io, weights=weights, spatial_methods=spatial_methods,
76+
physics_schemes=physics_schemes)
77+
# ------------------------------------------------------------------------ #
78+
# Initial conditions
79+
# ------------------------------------------------------------------------ #
80+
81+
xc_init = 0.25*L
82+
xc_end = 0.75*L
83+
umax = 0.5*L/tmax
84+
85+
# Get minimum distance on periodic interval to xc
86+
x_init = conditional(sqrt((x[0] - xc_init)**2) < 0.5*L,
87+
x[0] - xc_init, L + x[0] - xc_init)
88+
89+
x_end = conditional(sqrt((x[0] - xc_end)**2) < 0.5*L,
90+
x[0] - xc_end, L + x[0] - xc_end)
91+
92+
f_init = 5.0
93+
f_end = f_init / 2.0
94+
f_width_init = L / 10.0
95+
f_width_end = f_width_init * 2.0
96+
f_init_expr = f_init*exp(-(x_init / f_width_init)**2)
97+
98+
# The end Gaussian should be advected by half the domain
99+
# length, be more spread out due to the dissipation,
100+
# and includes more mass due to the source term.
101+
f_end_expr = 0.1 + f_end*exp(-(x_end / f_width_end)**2)
102+
103+
stepper.fields('f').interpolate(f_init_expr)
104+
stepper.fields('u').interpolate(as_vector([Constant(umax)]))
105+
f_end = stepper.fields('f_end', space=V)
106+
f_end.interpolate(f_end_expr)
107+
108+
# ------------------------------------------------------------------------ #
109+
# Run
110+
# ------------------------------------------------------------------------ #
111+
112+
stepper.run(0, tmax=tmax)
113+
114+
error = norm(stepper.fields('f') - f_end) / norm(f_end)
115+
116+
return error
117+
118+
119+
@pytest.mark.parametrize("timestepper", ["split1", "split2", "split3"])
120+
def test_split_timestepper_adv_diff_physics(tmpdir, timestepper):
121+
122+
tol = 0.015
123+
error = run_split_timestepper_adv_diff_physics(tmpdir, timestepper)
124+
print(error)
125+
assert error < tol, 'The split timestepper in the advection-diffusion' + \
126+
'equation with source physics has an error greater than ' + \
127+
'the permitted tolerance'

0 commit comments

Comments
 (0)