Neaten up subcycling options (#591)

Co-authored-by: Dr Jemma Shipton <[email protected]>
firedrakeproject · Dec 16, 2024 · 6f4150d · 6f4150d
1 parent 505e47f
commit 6f4150d
Show file tree

Hide file tree

Showing 10 changed files with 243 additions and 124 deletions.
diff --git a/examples/compressible_euler/dcmip_3_1_gravity_wave.py b/examples/compressible_euler/dcmip_3_1_gravity_wave.py
@@ -17,6 +17,7 @@
     TrapeziumRule, SUPGOptions, lonlatr_from_xyz, CompressibleParameters,
     CompressibleEulerEquations, CompressibleSolver, ZonalComponent,
     compressible_hydrostatic_balance, Perturbation, GeneralCubedSphereMesh,
+    SubcyclingOptions
 )
 
 dcmip_3_1_gravity_wave_defaults = {
@@ -96,10 +97,14 @@ def dcmip_3_1_gravity_wave(
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
+    subcycling_options = SubcyclingOptions(fixed_subcycles=2)
     transported_fields = [
         TrapeziumRule(domain, "u"),
-        SSPRK3(domain, "rho", fixed_subcycles=2),
-        SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2)
+        SSPRK3(domain, "rho", subcycling_options=subcycling_options),
+        SSPRK3(
+            domain, "theta", options=SUPGOptions(),
+            subcycling_options=subcycling_options
+        ),
     ]
     transport_methods = [
         DGUpwind(eqns, field) for field in ["u", "rho", "theta"]

diff --git a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py
@@ -22,7 +22,7 @@
     SUPGOptions, CourantNumber, Perturbation, Gradient,
     CompressibleParameters, CompressibleEulerEquations, CompressibleSolver,
     compressible_hydrostatic_balance, logger, RichardsonNumber,
-    RungeKuttaFormulation
+    RungeKuttaFormulation, SubcyclingOptions
 )
 
 skamarock_klemp_nonhydrostatic_defaults = {
@@ -106,10 +106,17 @@ def skamarock_klemp_nonhydrostatic(
 
     # Transport schemes
     theta_opts = SUPGOptions()
+    subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25)
     transported_fields = [
-        SSPRK3(domain, "u", subcycle_by_courant=0.25),
-        SSPRK3(domain, "rho", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear),
-        SSPRK3(domain, "theta", subcycle_by_courant=0.25, options=theta_opts)
+        SSPRK3(domain, "u", subcycling_options=subcycling_options),
+        SSPRK3(
+            domain, "rho", subcycling_options=subcycling_options,
+            rk_formulation=RungeKuttaFormulation.linear
+        ),
+        SSPRK3(
+            domain, "theta", subcycling_options=subcycling_options,
+            options=theta_opts
+        )
     ]
     transport_methods = [
         DGUpwind(eqns, "u"),

diff --git a/examples/shallow_water/thermal_williamson_2.py b/examples/shallow_water/thermal_williamson_2.py
@@ -18,7 +18,8 @@
     TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations,
     RelativeVorticity, PotentialVorticity, SteadyStateError,
     ZonalComponent, MeridionalComponent, ThermalSWSolver,
-    xyz_vector_from_lonlatr, lonlatr_from_xyz, GeneralIcosahedralSphereMesh
+    xyz_vector_from_lonlatr, lonlatr_from_xyz, GeneralIcosahedralSphereMesh,
+    SubcyclingOptions
 )
 
 thermal_williamson_2_defaults = {
@@ -88,10 +89,11 @@ def thermal_williamson_2(
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
+    subcycling_options = SubcyclingOptions(fixed_subcycles=2)
     transported_fields = [
         TrapeziumRule(domain, "u"),
-        SSPRK3(domain, "D", fixed_subcycles=2),
-        SSPRK3(domain, "b", fixed_subcycles=2)
+        SSPRK3(domain, "D", subcycling_options=subcycling_options),
+        SSPRK3(domain, "b", subcycling_options=subcycling_options)
     ]
     transport_methods = [
         DGUpwind(eqns, "u"),

diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py
@@ -16,7 +16,7 @@
     ShallowWaterKineticEnergy, ShallowWaterPotentialEnergy,
     ShallowWaterPotentialEnstrophy, rotated_lonlatr_coords,
     ZonalComponent, MeridionalComponent, rotated_lonlatr_vectors,
-    GeneralIcosahedralSphereMesh
+    GeneralIcosahedralSphereMesh, SubcyclingOptions
 )
 
 williamson_2_defaults = {
@@ -87,9 +87,10 @@ def williamson_2(
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
+    subcycling_options = SubcyclingOptions(fixed_subcycles=2)
     transported_fields = [
         TrapeziumRule(domain, "u"),
-        SSPRK3(domain, "D", fixed_subcycles=2)]
+        SSPRK3(domain, "D", subcycling_options=subcycling_options)]
     transport_methods = [
         DGUpwind(eqns, "u"),
         DGUpwind(eqns, "D")

diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py
@@ -14,7 +14,8 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     ShallowWaterParameters, ShallowWaterEquations, Sum,
     lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent,
-    MeridionalComponent, RelativeVorticity, RungeKuttaFormulation
+    MeridionalComponent, RelativeVorticity, RungeKuttaFormulation,
+    SubcyclingOptions
 )
 
 williamson_5_defaults = {
@@ -84,9 +85,13 @@ def williamson_5(
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
+    subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25)
     transported_fields = [
-        SSPRK3(domain, "u", subcycle_by_courant=0.25),
-        SSPRK3(domain, "D", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear)
+        SSPRK3(domain, "u", subcycling_options=subcycling_options),
+        SSPRK3(
+            domain, "D", subcycling_options=subcycling_options,
+            rk_formulation=RungeKuttaFormulation.linear
+        )
     ]
     transport_methods = [
         DGUpwind(eqns, "u"),

diff --git a/gusto/core/configuration.py b/gusto/core/configuration.py
@@ -10,7 +10,8 @@
     "ShallowWaterParameters",
     "EmbeddedDGOptions", "ConservativeEmbeddedDGOptions", "RecoveryOptions",
     "ConservativeRecoveryOptions", "SUPGOptions", "MixedFSOptions",
-    "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters"
+    "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters",
+    "SubcyclingOptions"
 ]
 
 
@@ -252,3 +253,31 @@ class BoundaryLayerParameters(Configuration):
     coeff_evap = 1.1e-3         # Dimensionless surface evaporation coefficient
     height_surface_layer = 75.  # Height (m) of surface level (usually lowest level)
     mu = 100.                   # Interior penalty coefficient for vertical diffusion
+
+
+class SubcyclingOptions(Configuration):
+    """
+    Describes the process of subcycling a time discretisation, by dividing the
+    time step into a number of smaller substeps.
+
+    NB: cannot provide both the fixed_subcycles and max_subcycles parameters,
+    which will raise an error.
+    """
+
+    # Either None, or an integer, giving the number of subcycles to take
+    fixed_subcycles = None
+
+    # If adaptive subcycling, the maximum number of subcycles to take
+    max_subcycles = 10
+
+    # Either None or a float, giving the maximum Courant number for one step
+    subcycle_by_courant = None
+
+    def check_options(self):
+        """Checks that the subcycling options are valid."""
+
+        if (self.fixed_subcycles is not None
+                and self.subcycle_by_courant is not None):
+            raise ValueError(
+                "Cannot provide both fixed_subcycles and subcycle_by_courant"
+                + "parameters.")
diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py
@@ -87,7 +87,7 @@ class ExplicitRungeKutta(ExplicitTimeDiscretisation):
     # ---------------------------------------------------------------------------
 
     def __init__(self, domain, butcher_matrix, field_name=None,
-                 fixed_subcycles=None, subcycle_by_courant=None,
+                 subcycling_options=None,
                  rk_formulation=RungeKuttaFormulation.increment,
                  solver_parameters=None, limiter=None, options=None):
         """
@@ -98,15 +98,9 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                 of a butcher tableau defining a given Runge Kutta scheme.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            fixed_subcycles (int, optional): the fixed number of sub-steps to
-                perform. This option cannot be specified with the
-                `subcycle_by_courant` argument. Defaults to None.
-            subcycle_by_courant (float, optional): specifying this option will
-                make the scheme perform adaptive sub-cycling based on the
-                Courant number. The specified argument is the maximum Courant
-                for one sub-cycle. Defaults to None, in which case adaptive
-                sub-cycling is not used. This option cannot be specified with
-                the `fixed_subcycles` argument.
+            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.
@@ -120,8 +114,7 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                 recovery method. Defaults to None.
         """
         super().__init__(domain, field_name=field_name,
-                         fixed_subcycles=fixed_subcycles,
-                         subcycle_by_courant=subcycle_by_courant,
+                         subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
         self.butcher_matrix = butcher_matrix
@@ -489,8 +482,7 @@ class ForwardEuler(ExplicitRungeKutta):
     y^(n+1) = y^n + dt*k0                                                     \n
     """
     def __init__(
-            self, domain, field_name=None,
-            fixed_subcycles=None, subcycle_by_courant=None,
+            self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
             solver_parameters=None, limiter=None, options=None
     ):
@@ -500,15 +492,9 @@ def __init__(
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            fixed_subcycles (int, optional): the fixed number of sub-steps to
-                perform. This option cannot be specified with the
-                `subcycle_by_courant` argument. Defaults to None.
-            subcycle_by_courant (float, optional): specifying this option will
-                make the scheme perform adaptive sub-cycling based on the
-                Courant number. The specified argument is the maximum Courant
-                for one sub-cycle. Defaults to None, in which case adaptive
-                sub-cycling is not used. This option cannot be specified with
-                the `fixed_subcycles` argument.
+            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.
@@ -525,8 +511,7 @@ def __init__(
         butcher_matrix = np.array([1.]).reshape(1, 1)
 
         super().__init__(domain, butcher_matrix, field_name=field_name,
-                         fixed_subcycles=fixed_subcycles,
-                         subcycle_by_courant=subcycle_by_courant,
+                         subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
@@ -543,8 +528,7 @@ class SSPRK3(ExplicitRungeKutta):
     y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2)                                 \n
     """
     def __init__(
-            self, domain, field_name=None,
-            fixed_subcycles=None, subcycle_by_courant=None,
+            self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
             solver_parameters=None, limiter=None, options=None
     ):
@@ -554,15 +538,9 @@ def __init__(
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            fixed_subcycles (int, optional): the fixed number of sub-steps to
-                perform. This option cannot be specified with the
-                `subcycle_by_courant` argument. Defaults to None.
-            subcycle_by_courant (float, optional): specifying this option will
-                make the scheme perform adaptive sub-cycling based on the
-                Courant number. The specified argument is the maximum Courant
-                for one sub-cycle. Defaults to None, in which case adaptive
-                sub-cycling is not used. This option cannot be specified with
-                the `fixed_subcycles` argument.
+            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.
@@ -582,8 +560,7 @@ def __init__(
             [1./6., 1./6., 2./3.]
         ])
         super().__init__(domain, butcher_matrix, field_name=field_name,
-                         fixed_subcycles=fixed_subcycles,
-                         subcycle_by_courant=subcycle_by_courant,
+                         subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
@@ -605,8 +582,7 @@ class RK4(ExplicitRungeKutta):
     where superscripts indicate the time-level.                               \n
     """
     def __init__(
-            self, domain, field_name=None,
-            fixed_subcycles=None, subcycle_by_courant=None,
+            self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
             solver_parameters=None, limiter=None, options=None
     ):
@@ -616,15 +592,9 @@ def __init__(
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            fixed_subcycles (int, optional): the fixed number of sub-steps to
-                perform. This option cannot be specified with the
-                `subcycle_by_courant` argument. Defaults to None.
-            subcycle_by_courant (float, optional): specifying this option will
-                make the scheme perform adaptive sub-cycling based on the
-                Courant number. The specified argument is the maximum Courant
-                for one sub-cycle. Defaults to None, in which case adaptive
-                sub-cycling is not used. This option cannot be specified with
-                the `fixed_subcycles` argument.
+            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.
@@ -644,8 +614,7 @@ def __init__(
             [1./6., 1./3., 1./3., 1./6.]
         ])
         super().__init__(domain, butcher_matrix, field_name=field_name,
-                         fixed_subcycles=fixed_subcycles,
-                         subcycle_by_courant=subcycle_by_courant,
+                         subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)
@@ -665,8 +634,7 @@ class Heun(ExplicitRungeKutta):
     number.
     """
     def __init__(
-            self, domain, field_name=None,
-            fixed_subcycles=None, subcycle_by_courant=None,
+            self, domain, field_name=None, subcycling_options=None,
             rk_formulation=RungeKuttaFormulation.increment,
             solver_parameters=None, limiter=None, options=None
     ):
@@ -676,15 +644,9 @@ def __init__(
                 mesh and the compatible function spaces.
             field_name (str, optional): name of the field to be evolved.
                 Defaults to None.
-            fixed_subcycles (int, optional): the fixed number of sub-steps to
-                perform. This option cannot be specified with the
-                `subcycle_by_courant` argument. Defaults to None.
-            subcycle_by_courant (float, optional): specifying this option will
-                make the scheme perform adaptive sub-cycling based on the
-                Courant number. The specified argument is the maximum Courant
-                for one sub-cycle. Defaults to None, in which case adaptive
-                sub-cycling is not used. This option cannot be specified with the
-                `fixed_subcycles` argument.
+            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.
@@ -703,8 +665,7 @@ def __init__(
             [0.5, 0.5]
         ])
         super().__init__(domain, butcher_matrix, field_name=field_name,
-                         fixed_subcycles=fixed_subcycles,
-                         subcycle_by_courant=subcycle_by_courant,
+                         subcycling_options=subcycling_options,
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options)