From d61f81f4a0e111ab3e492b4a78349609595f72cc Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 19 Dec 2024 16:50:43 +0000 Subject: [PATCH] Add SSPRK schemes with higher stage count (#606) Co-authored-by: Thomas Bendall --- .../explicit_runge_kutta.py | 240 +++++++++++++++--- .../model/test_time_discretisation.py | 60 ++++- 2 files changed, 256 insertions(+), 44 deletions(-) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index b167cf80a..16732bd50 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -15,8 +15,8 @@ __all__ = [ - "ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun", - "RungeKuttaFormulation" + "ForwardEuler", "ExplicitRungeKutta", "SSPRK2", "SSPRK3", "SSPRK4", + "RK4", "Heun", "RungeKuttaFormulation" ] @@ -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: @@ -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, @@ -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: @@ -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, @@ -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, @@ -686,10 +826,11 @@ 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, @@ -697,3 +838,24 @@ def __init__( 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) diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 6d484bfd1..0c259e221 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -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): @@ -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":