Skip to content

Commit

Permalink
Add SSPRK schemes with higher stage count (#606)
Browse files Browse the repository at this point in the history 
Co-authored-by: Thomas Bendall <[email protected]>
  • Loading branch information
@JHopeCollins @tommbendall
JHopeCollins and tommbendall authored Dec 19, 2024
1 parent e305589 commit d61f81f
Show file tree
Hide file tree
Showing 2 changed files with 256 additions and 44 deletions.
240 changes: 201 additions & 39 deletions gusto/time_discretisation/explicit_runge_kutta.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@


__all__ = [
"ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun",
"RungeKuttaFormulation"
"ForwardEuler", "ExplicitRungeKutta", "SSPRK2", "SSPRK3", "SSPRK4",
"RK4", "Heun", "RungeKuttaFormulation"
]


Expand Down Expand Up @@ -530,21 +530,132 @@ def __init__(
augmentation=augmentation)


class SSPRK2(ExplicitRungeKutta):
u"""
Implements 2nd order Strong-Stability-Preserving Runge-Kutta methods
for solving ∂y/∂t = F(y). \n
Spiteri & Ruuth, 2002, SIAM J. Numer. Anal. \n
"A new class of optimal high-order strong-stability-preserving time \n
discretisation methods". \n
The 2-stage method (Heun's method) can be written as: \n
k0 = F[y^n] \n
k1 = F{y^n + dt*k0] \n
y^(n+1) = y^n + (1/2)*dt*(k0+k1) \n
The 3-stage method can be written as: \n
k0 = F[y^n] \n
k1 = F[y^n + (1/2*dt*k0] \n
k2 = F[y^n + (1/2)*dt*(k0+k1)] \n
y^(n+1) = y^n + (1/3)*dt*(k0 + k1 + k2) \n
The 4-stage method can be written as: \n
k0 = F[y^n] \n
k1 = F[y^n + (1/3)*dt*k1] \n
k2 = F[y^n + (1/3)*dt*(k0+k1)] \n
k3 = F[y^n + (1/3)*dt*(k0+k1+k2)] \n
y^(n+1) = y^n + (1/4)*dt*(k0 + k1 + k2 + k3) \n
"""
def __init__(
self, domain, field_name=None, subcycling_options=None,
rk_formulation=RungeKuttaFormulation.increment,
solver_parameters=None, limiter=None, options=None,
augmentation=None, stages=2
):
"""
Args:
domain (:class:`Domain`): the model's domain object, containing the
mesh and the compatible function spaces.
field_name (str, optional): name of the field to be evolved.
Defaults to None.
subcycling_options(:class:`SubcyclingOptions`, optional): an object
containing options for subcycling the time discretisation.
Defaults to None.
rk_formulation (:class:`RungeKuttaFormulation`, optional):
an enumerator object, describing the formulation of the Runge-
Kutta scheme. Defaults to the increment form.
solver_parameters (dict, optional): dictionary of parameters to
pass to the underlying solver. Defaults to None.
limiter (:class:`Limiter` object, optional): a limiter to apply to
the evolving field to enforce monotonicity. Defaults to None.
options (:class:`AdvectionOptions`, optional): an object containing
options to either be passed to the spatial discretisation, or
to control the "wrapper" methods, such as Embedded DG or a
recovery method. Defaults to None.
augmentation (:class:`Augmentation`): allows the equation solved in
this time discretisation to be augmented, for instances with
extra terms of another auxiliary variable. Defaults to None.
stages (int, optional): number of stages: (2, 3, 4). Defaults to 2.
"""

if stages == 2:
butcher_matrix = np.array([
[1., 0.],
[0.5, 0.5]
])
self.cfl_limit = 1

elif stages == 3:
butcher_matrix = np.array([
[1./2., 0., 0.],
[1./2., 1./2., 0.],
[1./3., 1./3., 1./3.]
])
self.cfl_limit = 2
elif stages == 4:
butcher_matrix = np.array([
[1./3., 0., 0., 0.],
[1./3., 1./3., 0., 0.],
[1./3., 1./3., 1./3., 0.],
[1./4., 1./4., 1./4., 1./4.]
])
self.cfl_limit = 3
else:
raise ValueError(f"{stages} stage 2rd order SSPRK not implemented")

super().__init__(domain, butcher_matrix, field_name=field_name,
subcycling_options=subcycling_options,
rk_formulation=rk_formulation,
solver_parameters=solver_parameters,
limiter=limiter, options=options,
augmentation=augmentation)


class SSPRK3(ExplicitRungeKutta):
u"""
Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method
for solving ∂y/∂t = F(y). It can be written as: \n
Implements 3rd order Strong-Stability-Preserving Runge-Kutta methods
for solving ∂y/∂t = F(y). \n
Spiteri & Ruuth, 2002, SIAM J. Numer. Anal. \n
"A new class of optimal high-order strong-stability-preserving time \n
discretisation methods". \n
The 3-stage method can be written as: \n
k0 = F[y^n] \n
k1 = F[y^n + dt*k1] \n
k2 = F[y^n + (1/4)*dt*(k0+k1)] \n
y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) \n
The 4-stage method can be written as: \n
k0 = F[y^n] \n
k1 = F[y^n + (1/2)*dt*k1] \n
k2 = F[y^n + (1/2)*dt*(k0+k1)] \n
k3 = F[y^n + (1/6)*dt*(k0+k1+k2)] \n
y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + k2 + 3*k3) \n
The 5-stage method has numerically optimised coefficients. \n
"""
def __init__(
self, domain, field_name=None, subcycling_options=None,
rk_formulation=RungeKuttaFormulation.increment,
solver_parameters=None, limiter=None, options=None,
augmentation=None
augmentation=None, stages=3
):
"""
Args:
Expand All @@ -569,13 +680,36 @@ def __init__(
augmentation (:class:`Augmentation`): allows the equation solved in
this time discretisation to be augmented, for instances with
extra terms of another auxiliary variable. Defaults to None.
stages (int, optional): number of stages: (3, 4, 5). Defaults to 3.
"""

butcher_matrix = np.array([
[1., 0., 0.],
[1./4., 1./4., 0.],
[1./6., 1./6., 2./3.]
])
if stages == 3:
butcher_matrix = np.array([
[1., 0., 0.],
[1./4., 1./4., 0.],
[1./6., 1./6., 2./3.]
])
self.cfl_limit = 1
elif stages == 4:
butcher_matrix = np.array([
[1./2., 0., 0., 0.],
[1./2., 1./2., 0., 0.],
[1./6., 1./6., 1./6., 0.],
[1./6., 1./6., 1./6., 1./2.]
])
self.cfl_limit = 2
elif stages == 5:
self.cfl_limit = 2.65062919294483
butcher_matrix = np.array([
[0.37726891511710, 0., 0., 0., 0.],
[0.37726891511710, 0.37726891511710, 0., 0., 0.],
[0.16352294089771, 0.16352294089771, 0.16352294089771, 0., 0.],
[0.14904059394856, 0.14831273384724, 0.14831273384724, 0.34217696850008, 0.],
[0.19707596384481, 0.11780316509765, 0.11709725193772, 0.27015874934251, 0.29786487010104]
])
else:
raise ValueError(f"{stages} stage 3rd order SSPRK not implemented")

super().__init__(domain, butcher_matrix, field_name=field_name,
subcycling_options=subcycling_options,
rk_formulation=rk_formulation,
Expand All @@ -584,26 +718,22 @@ def __init__(
augmentation=augmentation)


class RK4(ExplicitRungeKutta):
class SSPRK4(ExplicitRungeKutta):
u"""
Implements the classic 4-stage Runge-Kutta method.
Implements 4th order Strong-Stability-Preserving Runge-Kutta methods
for solving ∂y/∂t = F(y). \n
The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be
written as: \n
Spiteri & Ruuth, 2002, SIAM J. Numer. Anal. \n
"A new class of optimal high-order strong-stability-preserving time \n
discretisation methods". \n
k0 = F[y^n] \n
k1 = F[y^n + 1/2*dt*k1] \n
k2 = F[y^n + 1/2*dt*k2] \n
k3 = F[y^n + dt*k3] \n
y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3) \n
where superscripts indicate the time-level. \n
The 5-stage method has numerically optimised coefficients. \n
"""
def __init__(
self, domain, field_name=None, subcycling_options=None,
rk_formulation=RungeKuttaFormulation.increment,
solver_parameters=None, limiter=None, options=None,
augmentation=None
augmentation=None, stages=5
):
"""
Args:
Expand All @@ -628,13 +758,21 @@ def __init__(
augmentation (:class:`Augmentation`): allows the equation solved in
this time discretisation to be augmented, for instances with
extra terms of another auxiliary variable. Defaults to None.
stages (int, optional): number of stages: (5,). Defaults to 5.
"""
butcher_matrix = np.array([
[0.5, 0., 0., 0.],
[0., 0.5, 0., 0.],
[0., 0., 1., 0.],
[1./6., 1./3., 1./3., 1./6.]
])

if stages == 5:
self.cfl_limit = 1.50818004975927
butcher_matrix = np.array([
[0.39175222700392, 0., 0., 0., 0.],
[0.21766909633821, 0.36841059262959, 0., 0., 0.],
[0.08269208670950, 0.13995850206999, 0.25189177424738, 0., 0.],
[0.06796628370320, 0.11503469844438, 0.20703489864929, 0.54497475021237, 0.],
[0.14681187618661, 0.24848290924556, 0.10425883036650, 0.27443890091960, 0.22600748319395]
])
else:
raise ValueError(f"{stages} stage 4rd order SSPRK not implemented")

super().__init__(domain, butcher_matrix, field_name=field_name,
subcycling_options=subcycling_options,
rk_formulation=rk_formulation,
Expand All @@ -643,18 +781,20 @@ def __init__(
augmentation=augmentation)


class Heun(ExplicitRungeKutta):
class RK4(ExplicitRungeKutta):
u"""
Implements Heun's method.
Implements the classic 4-stage Runge-Kutta method.
The 2-stage Runge-Kutta scheme known as Heun's method,for solving
∂y/∂t = F(y). It can be written as: \n
The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be
written as: \n
y_1 = F[y^n] \n
y^(n+1) = (1/2)y^n + (1/2)F[y_1] \n
k0 = F[y^n] \n
k1 = F[y^n + 1/2*dt*k1] \n
k2 = F[y^n + 1/2*dt*k2] \n
k3 = F[y^n + dt*k3] \n
y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3) \n
where superscripts indicate the time-level and subscripts indicate the stage
number.
where superscripts indicate the time-level. \n
"""
def __init__(
self, domain, field_name=None, subcycling_options=None,
Expand Down Expand Up @@ -686,14 +826,36 @@ def __init__(
this time discretisation to be augmented, for instances with
extra terms of another auxiliary variable. Defaults to None.
"""

butcher_matrix = np.array([
[1., 0.],
[0.5, 0.5]
[0.5, 0., 0., 0.],
[0., 0.5, 0., 0.],
[0., 0., 1., 0.],
[1./6., 1./3., 1./3., 1./6.]
])
super().__init__(domain, butcher_matrix, field_name=field_name,
subcycling_options=subcycling_options,
rk_formulation=rk_formulation,
solver_parameters=solver_parameters,
limiter=limiter, options=options,
augmentation=augmentation)


def Heun(domain, field_name=None, subcycling_options=None,
rk_formulation=RungeKuttaFormulation.increment,
solver_parameters=None, limiter=None, options=None,
augmentation=None):
u"""
Implements Heun's method.
The 2-stage Runge-Kutta scheme known as Heun's method,for solving
∂y/∂t = F(y). It can be written as: \n
y_1 = F[y^n] \n
y^(n+1) = (1/2)y^n + (1/2)F[y_1] \n
where superscripts indicate the time-level and subscripts indicate the stage
number.
"""
return SSPRK2(domain, field_name=field_name, subcycling_options=subcycling_options,
rk_formulation=rk_formulation, solver_parameters=solver_parameters,
limiter=limiter, options=options, augmentation=augmentation, stages=2)
60 changes: 55 additions & 5 deletions integration-tests/model/test_time_discretisation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,19 @@ def run(timestepper, tmax, f_end):

@pytest.mark.parametrize(
"scheme", [
"ssprk3_increment", "TrapeziumRule", "ImplicitMidpoint", "QinZhang",
"ssprk2_increment_2", "ssprk2_predictor_2", "ssprk2_linear_2",
"ssprk2_increment_3", "ssprk2_predictor_3", "ssprk2_linear_3",
"ssprk2_increment_4", "ssprk2_predictor_4", "ssprk2_linear_4",
"ssprk3_increment_3", "ssprk3_predictor_3", "ssprk3_linear_3",
"ssprk3_increment_4", "ssprk3_predictor_4", "ssprk3_linear_4",
"ssprk3_increment_5", "ssprk3_predictor_5", "ssprk3_linear_5",
"ssprk4_increment_5", "ssprk4_predictor_5", "ssprk4_linear_5",
"TrapeziumRule", "ImplicitMidpoint", "QinZhang",
"RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog",
"AdamsMoulton", "AdamsMoulton", "ssprk3_predictor", "ssprk3_linear"
"AdamsMoulton", "AdamsMoulton"
]
)
def test_time_discretisation(tmpdir, scheme, tracer_setup):
Expand All @@ -30,12 +40,52 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup):
V = domain.spaces("DG")
eqn = AdvectionEquation(domain, V, "f")

if scheme == "ssprk3_increment":
if scheme == "ssprk2_increment_2":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment)
elif scheme == "ssprk2_predictor_2":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor)
elif scheme == "ssprk2_linear_2":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear)
elif scheme == "ssprk2_increment_3":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment, stages=3)
elif scheme == "ssprk2_predictor_3":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=3)
elif scheme == "ssprk2_linear_3":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear, stages=3)
elif scheme == "ssprk2_increment_4":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.increment, stages=4)
elif scheme == "ssprk2_predictor_4":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=4)
elif scheme == "ssprk2_linear_4":
transport_scheme = SSPRK2(domain, rk_formulation=RungeKuttaFormulation.linear, stages=4)

elif scheme == "ssprk3_increment_3":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment)
elif scheme == "ssprk3_predictor":
elif scheme == "ssprk3_predictor_3":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor)
elif scheme == "ssprk3_linear":
elif scheme == "ssprk3_linear_3":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear)
elif scheme == "ssprk3_increment_4":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment, stages=4)
elif scheme == "ssprk3_predictor_4":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=4)
elif scheme == "ssprk3_linear_4":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear, stages=4)

elif scheme == "ssprk3_increment_5":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment, stages=5)
elif scheme == "ssprk3_predictor_5":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor, stages=5)
elif scheme == "ssprk3_linear_5":
transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear, stages=5)

elif scheme == "ssprk4_increment_5":
transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.increment)
elif scheme == "ssprk4_predictor_5":
transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.predictor)
elif scheme == "ssprk4_linear_5":
transport_scheme = SSPRK4(domain, rk_formulation=RungeKuttaFormulation.linear)

elif scheme == "TrapeziumRule":
transport_scheme = TrapeziumRule(domain)
elif scheme == "ImplicitMidpoint":
Expand Down

0 comments on commit d61f81f

Please sign in to comment.