Advective-then-Flux form for Runge-Kutta schemes

examples/compressible_euler/skamarock_klemp_nonhydrostatic.py

@@ -21,7 +21,8 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, SUPGOptions, CourantNumber, Perturbation, Gradient,
     CompressibleParameters, CompressibleEulerEquations, CompressibleSolver,
-    compressible_hydrostatic_balance, logger, RichardsonNumber
+    compressible_hydrostatic_balance, logger, RichardsonNumber,
+    RungeKuttaFormulation
 )
 
 skamarock_klemp_nonhydrostatic_defaults = {
@@ -106,13 +107,13 @@ def skamarock_klemp_nonhydrostatic(
     # Transport schemes
     theta_opts = SUPGOptions()
     transported_fields = [
-        TrapeziumRule(domain, "u"),
-        SSPRK3(domain, "rho"),
-        SSPRK3(domain, "theta", options=theta_opts)
+        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)
     ]
     transport_methods = [
         DGUpwind(eqns, "u"),
-        DGUpwind(eqns, "rho"),
+        DGUpwind(eqns, "rho", advective_then_flux=True),
         DGUpwind(eqns, "theta", ibp=theta_opts.ibp)
     ]
 

examples/shallow_water/moist_thermal_williamson_5.py

@@ -24,7 +24,8 @@
     ShallowWaterParameters, ShallowWaterEquations, Sum,
     lonlatr_from_xyz, InstantRain, SWSaturationAdjustment, WaterVapour,
     CloudWater, Rain, GeneralIcosahedralSphereMesh, RelativeVorticity,
-    ZonalComponent, MeridionalComponent
+    ZonalComponent, MeridionalComponent, RungeKuttaFormulation, SSPRK3,
+    SemiImplicitQuasiNewton, ThermalSWSolver
 )
 
 moist_thermal_williamson_5_defaults = {
@@ -142,13 +143,28 @@ def gamma_v(x_in):
         gamma_r=gamma_r
     )
 
+    transported_fields = [
+        SSPRK3(domain, "u", subcycle_by_courant=0.25),
+        SSPRK3(domain, "D", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear),
+        SSPRK3(domain, "b", subcycle_by_courant=0.25),
+        SSPRK3(domain, "water_vapour", subcycle_by_courant=0.25),
+        SSPRK3(domain, "cloud_water", subcycle_by_courant=0.25),
+    ]
     transport_methods = [
-        DGUpwind(eqns, field_name) for field_name in eqns.field_names
+        DGUpwind(eqns, "u"),
+        DGUpwind(eqns, "D", advective_then_flux=True),
+        DGUpwind(eqns, "b"),
+        DGUpwind(eqns, "water_vapour"),
+        DGUpwind(eqns, "cloud_water")
     ]
 
-    # Timestepper
-    stepper = Timestepper(
-        eqns, RK4(domain), io, spatial_methods=transport_methods
+    # Linear solver
+    linear_solver = ThermalSWSolver(eqns)
+
+    # Time stepper
+    stepper = SemiImplicitQuasiNewton(
+        eqns, io, transported_fields, transport_methods,
+        linear_solver=linear_solver
     )
 
     # ------------------------------------------------------------------------ #

examples/shallow_water/williamson_5.py

@@ -14,7 +14,7 @@
     Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind,
     TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, Sum,
     lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent,
-    MeridionalComponent, RelativeVorticity
+    MeridionalComponent, RelativeVorticity, RungeKuttaFormulation
 )
 
 williamson_5_defaults = {
@@ -83,8 +83,14 @@ def williamson_5(
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
-    transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")]
-    transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
+    transported_fields = [
+        SSPRK3(domain, "u", subcycle_by_courant=0.25),
+        SSPRK3(domain, "D", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear)
+    ]
+    transport_methods = [
+        DGUpwind(eqns, "u"),
+        DGUpwind(eqns, "D", advective_then_flux=True)
+    ]
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(

gusto/core/labels.py

@@ -97,6 +97,8 @@ def __call__(self, target, value=None):
 linearisation = Label("linearisation", validator=lambda value: type(value) in [LabelledForm, Term])
 mass_weighted = Label("mass_weighted", validator=lambda value: type(value) in [LabelledForm, Term])
 ibp_label = Label("ibp", validator=lambda value: type(value) == IntegrateByParts)
+all_but_last = Label("all_but_last", validator=lambda value: type(value) in [LabelledForm, Term])
+
 
 # labels for terms in the equations
 time_derivative = Label("time_derivative")

gusto/spatial_methods/transport_methods.py

@@ -8,8 +8,10 @@
 )
 from firedrake.fml import Term, keep, drop
 from gusto.core.configuration import IntegrateByParts, TransportEquationType
-from gusto.core.labels import (prognostic, transport, transporting_velocity, ibp_label,
-                               mass_weighted)
+from gusto.core.labels import (
+    prognostic, transport, transporting_velocity, ibp_label, mass_weighted,
+    all_but_last
+)
 from gusto.core.logging import logger
 from gusto.spatial_methods.spatial_methods import SpatialMethod
 
@@ -83,6 +85,10 @@ def replace_form(self, equation):
             # Create new term
             new_term = Term(self.form.form, original_term.labels)
 
+            # Add all_but_last form
+            if hasattr(self, "all_but_last_form"):
+                new_term = all_but_last(new_term, self.all_but_last_form)
+
             # Check if this is a conservative transport
             if original_term.has_label(mass_weighted):
                 # Extract the original and discretised mass_weighted terms
@@ -151,7 +157,8 @@ class DGUpwind(TransportMethod):
     transported variable at facets.
     """
     def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE,
-                 vector_manifold_correction=False, outflow=False):
+                 vector_manifold_correction=False, outflow=False,
+                 advective_then_flux=False):
         """
         Args:
             equation (:class:`PrognosticEquation`): the equation, which includes
@@ -163,50 +170,82 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE,
                 vector manifold correction term. Defaults to False.
             outflow (bool, optional): whether to include outflow at the domain
                 boundaries, through exterior facet terms. Defaults to False.
+            advective_then_flux (bool, optional): whether to use the advective-
+                then-flux formulation. This uses the advective form of the
+                transport equation for all but the last steps of some
+                (potentially subcycled) Runge-Kutta scheme, before using the
+                conservative form for the final step to deliver a mass-
+                conserving increment. This optiona only makes sense to use with
+                Runge-Kutta, and should be used with the "linear" Runge-Kutta
+                formulation. Defaults to False, in which case the conservative
+                form is used for every step.
         """
 
         super().__init__(equation, variable)
         self.ibp = ibp
         self.vector_manifold_correction = vector_manifold_correction
         self.outflow = outflow
 
+        if (advective_then_flux
+                and self.transport_equation_type != TransportEquationType.conservative):
+            raise ValueError(
+                'DG Upwind: advective_then_flux form can only be used with '
+                + 'the conservative form of the transport equation'
+            )
+
         # -------------------------------------------------------------------- #
         # Determine appropriate form to use
         # -------------------------------------------------------------------- #
         # first check for 1d mesh and scalar velocity space
         if equation.domain.mesh.topological_dimension() == 1 and len(equation.domain.spaces("HDiv").shape) == 0:
             assert not vector_manifold_correction
             if self.transport_equation_type == TransportEquationType.advective:
-                form = upwind_advection_form_1d(self.domain, self.test,
-                                                self.field,
-                                                ibp=ibp, outflow=outflow)
+                form = upwind_advection_form_1d(
+                    self.domain, self.test, self.field, ibp=ibp,
+                    outflow=outflow
+                )
             elif self.transport_equation_type == TransportEquationType.conservative:
-                form = upwind_continuity_form_1d(self.domain, self.test,
-                                                 self.field,
-                                                 ibp=ibp, outflow=outflow)
+                form = upwind_continuity_form_1d(
+                    self.domain, self.test, self.field, ibp=ibp,
+                    outflow=outflow
+                )
 
         else:
             if self.transport_equation_type == TransportEquationType.advective:
                 if vector_manifold_correction:
-                    form = vector_manifold_advection_form(self.domain,
-                                                          self.test,
-                                                          self.field, ibp=ibp,
-                                                          outflow=outflow)
+                    form = vector_manifold_advection_form(
+                        self.domain, self.test, self.field, ibp=ibp,
+                        outflow=outflow
+                    )
                 else:
-                    form = upwind_advection_form(self.domain, self.test,
-                                                 self.field,
-                                                 ibp=ibp, outflow=outflow)
+                    form = upwind_advection_form(
+                        self.domain, self.test, self.field, ibp=ibp,
+                        outflow=outflow
+                    )
 
             elif self.transport_equation_type == TransportEquationType.conservative:
                 if vector_manifold_correction:
-                    form = vector_manifold_continuity_form(self.domain,
-                                                           self.test,
-                                                           self.field, ibp=ibp,
-                                                           outflow=outflow)
+                    form = vector_manifold_continuity_form(
+                        self.domain, self.test, self.field, ibp=ibp,
+                        outflow=outflow
+                    )
                 else:
-                    form = upwind_continuity_form(self.domain, self.test,
-                                                  self.field,
-                                                  ibp=ibp, outflow=outflow)
+                    form = upwind_continuity_form(
+                        self.domain, self.test, self.field, ibp=ibp,
+                        outflow=outflow
+                    )
+
+                if advective_then_flux and vector_manifold_correction:
+                    self.all_but_last_form = vector_manifold_advection_form(
+                        self.domain, self.test, self.field, ibp=ibp,
+                        outflow=outflow
+                    )
+
+                elif advective_then_flux:
+                    self.all_but_last_form = upwind_advection_form(
+                        self.domain, self.test, self.field, ibp=ibp,
+                        outflow=outflow
+                    )
 
             elif self.transport_equation_type == TransportEquationType.circulation:
                 if outflow: