Merge #1257

1257: Add SplitExplicitStepper with decoupled spin-down test r=blallen a=blallen

# Description

This PR adds the MultistateMultirateRungeKutta solver from the branches `bla/multistate_multirate` and `jmc/split_explicit_dvlp` in a new form as the SplitExplicitStepper. A simple test of running the `HBModel` and `SWModel` in parallel performing the hydrostatic spindown test. Both models converge to the analytic solution. A separate PR will be added that couples the two models with this test.

Originally, I attempted to include a working version of the `add_fast_steps` option from `jmc/split_explicit_dvlp`; however, that caused the `SWModel` to not converge properly. This feature will be revisited in a future PR. 



Co-authored-by: Brandon Allen <ballen@mit.edu>

docs/src/APIs/Numerics/ODESolvers/ODESolvers.md

@@ -44,6 +44,12 @@ ODESolvers.LSRKEulerMethod
 ODESolvers.MultirateRungeKutta
 ```
 
+## Split-explicit methods
+
+```@docs
+ODESolvers.SplitExplicitSolver
+```
+
 ## GARK methods
 
 ```@docs

slurmci-test.toml

@@ -39,7 +39,8 @@ cpu_gpu = [
   { file = "test/Numerics/Mesh/interpolation.jl", slurmargs = ["--ntasks=3"], args = [] },
   { file = "test/Ocean/ShallowWater/GyreDriver.jl", slurmargs = ["--ntasks=1"], args = [] },
   { file = "test/Ocean/ShallowWater/2D_hydrostatic_spindown.jl", slurmargs = ["--ntasks=1"], args = [] },
-  { file = "test/Ocean/HydrostaticBoussinesq/3D_hydrostatic_spindown.jl", slurmargs = ["--ntasks=3"], args = [] },
+  { file = "test/Ocean/HydrostaticBoussinesq/3D_hydrostatic_spindown.jl", slurmargs = ["--ntasks=1"], args = [] },
+  { file = "test/Ocean/SplitExplicit/hydrostatic_spindown.jl", slurmargs = ["--ntasks=1"], args = [] },
   { file = "tutorials/Numerics/DGMethods/nonnegative.jl", slurmargs = ["--ntasks=3"], args = [] },
   { file = "test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl", slurmargs = ["--ntasks=3"], args = [] },
   { file = "test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl", slurmargs = ["--ntasks=3"], args = [] },

src/BalanceLaws/interface.jl

@@ -326,3 +326,9 @@ num_hyperdiffusive(m::BalanceLaw, FT) = varsize(vars_hyperdiffusive(m, FT))
 num_integrals(m::BalanceLaw, FT) = varsize(vars_integrals(m, FT))
 num_reverse_integrals(m::BalanceLaw, FT) =
     varsize(vars_reverse_integrals(m, FT))
+
+### split explicit functions
+function initialize_states! end
+function tendency_from_slow_to_fast! end
+function cummulate_fast_solution! end
+function reconcile_from_fast_to_slow! end

src/ClimateMachine.jl

@@ -32,6 +32,7 @@ include(joinpath(
     "HydrostaticBoussinesq",
     "HydrostaticBoussinesqModel.jl",
 ))
+include(joinpath("Ocean", "SplitExplicit", "SplitExplicitModel.jl"))
 include(joinpath("Numerics", "SystemSolvers", "SystemSolvers.jl"))
 include(joinpath("Numerics", "ODESolvers", "ODESolvers.jl"))
 include(joinpath("Numerics", "ODESolvers", "GenericCallbacks.jl"))

src/Numerics/DGMethods/DGModel.jl

@@ -14,6 +14,7 @@ struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD}
     diffusion_direction::DD
     modeldata::MD
 end
+
 function DGModel(
     balance_law,
     grid,

src/Numerics/ODESolvers/ODESolvers.jl

@@ -165,5 +165,6 @@ include("StrongStabilityPreservingRungeKuttaMethod.jl")
 include("AdditiveRungeKuttaMethod.jl")
 include("MultirateInfinitesimalStepMethod.jl")
 include("MultirateRungeKuttaMethod.jl")
+include("SplitExplicitMethod.jl")
 
 end # module

src/Numerics/ODESolvers/SplitExplicitMethod.jl

@@ -0,0 +1,190 @@
+export SplitExplicitSolver
+
+using ..BalanceLaws:
+    initialize_states!,
+    tendency_from_slow_to_fast!,
+    cummulate_fast_solution!,
+    reconcile_from_fast_to_slow!
+
+LSRK2N = LowStorageRungeKutta2N
+
+@doc """
+    SplitExplicitSolver(slow_solver, fast_solver; dt, t0 = 0, coupled = true)
+
+This is a time stepping object for explicitly time stepping the differential
+equation given by the right-hand-side function `f` with the state `Q`, i.e.,
+
+```math
+  \\dot{Q_fast} = f_fast(Q_fast, Q_slow, t)
+  \\dot{Q_slow} = f_slow(Q_slow, Q_fast, t)
+```
+
+with the required time step size `dt` and optional initial time `t0`.  This
+time stepping object is intended to be passed to the `solve!` command.
+
+This method performs an operator splitting to timestep the vertical average
+of the model at a faster rate than the full model. This results in a first-
+order time stepper.
+
+""" SplitExplicitSolver
+mutable struct SplitExplicitSolver{SS, FS, RT, MSA} <: AbstractODESolver
+    "slow solver"
+    slow_solver::SS
+    "fast solver"
+    fast_solver::FS
+    "time step"
+    dt::RT
+    "time"
+    t::RT
+    "flag to couple the models"
+    coupled
+    "storage for transfer tendency"
+    dQ2fast::MSA
+
+    function SplitExplicitSolver(
+        slow_solver::LSRK2N,
+        fast_solver,
+        Q = nothing;
+        dt = getdt(slow_solver),
+        t0 = slow_solver.t,
+        coupled = true,
+    ) where {AT <: AbstractArray}
+        SS = typeof(slow_solver)
+        FS = typeof(fast_solver)
+        RT = real(eltype(slow_solver.dQ))
+
+        dQ2fast = similar(slow_solver.dQ)
+        MSA = typeof(dQ2fast)
+        return new{SS, FS, RT, MSA}(
+            slow_solver,
+            fast_solver,
+            RT(dt),
+            RT(t0),
+            coupled,
+            dQ2fast,
+        )
+    end
+end
+
+function dostep!(
+    Qvec,
+    split::SplitExplicitSolver{SS},
+    param,
+    time::Real,
+) where {SS <: LSRK2N}
+    slow = split.slow_solver
+    fast = split.fast_solver
+
+    Qslow = Qvec.slow
+    Qfast = Qvec.fast
+
+    dQslow = slow.dQ
+    dQ2fast = split.dQ2fast
+
+    slow_bl = slow.rhs!.balance_law
+    fast_bl = fast.rhs!.balance_law
+
+    groupsize = 256
+
+    slow_dt = getdt(slow)
+    fast_dt_in = getdt(fast)
+
+    for slow_s in 1:length(slow.RKA)
+        # Current slow state time
+        slow_stage_time = time + slow.RKC[slow_s] * slow_dt
+
+        # Initialize fast model and tendency adjustment
+        # before evalution of slow mode
+        if split.coupled
+            initialize_states!(
+                slow_bl,
+                fast_bl,
+                slow.rhs!,
+                fast.rhs!,
+                Qslow,
+                Qfast,
+            )
+
+            # Evaluate the slow mode
+            # --> save tendency for the fast
+            slow.rhs!(dQ2fast, Qslow, param, slow_stage_time, increment = false)
+
+            # vertically integrate slow tendency to advance fast equation
+            # and use vertical mean for slow model (negative source)
+            # ---> work with dQ2fast as input
+            tendency_from_slow_to_fast!(
+                slow_bl,
+                fast_bl,
+                slow.rhs!,
+                fast.rhs!,
+                Qslow,
+                Qfast,
+                dQ2fast,
+            )
+        end
+
+        # Compute (and RK update) slow tendency
+        slow.rhs!(dQslow, Qslow, param, slow_stage_time, increment = true)
+
+        # Update (RK-stage) slow state
+        event = Event(array_device(Qslow))
+        event = update!(array_device(Qslow), groupsize)(
+            realview(dQslow),
+            realview(Qslow),
+            slow.RKA[slow_s % length(slow.RKA) + 1],
+            slow.RKB[slow_s],
+            slow_dt,
+            nothing,
+            nothing,
+            nothing;
+            ndrange = length(realview(Qslow)),
+            dependencies = (event,),
+        )
+        wait(array_device(Qslow), event)
+
+        # Fractional time for slow stage
+        if slow_s == length(slow.RKA)
+            γ = 1 - slow.RKC[slow_s]
+        else
+            γ = slow.RKC[slow_s + 1] - slow.RKC[slow_s]
+        end
+
+        # RKB for the slow with fractional time factor remove (since full
+        # integration of fast will result in scaling by γ)
+        nsubsteps = fast_dt_in > 0 ? ceil(Int, γ * slow_dt / fast_dt_in) : 1
+        fast_dt = γ * slow_dt / nsubsteps
+
+        updatedt!(fast, fast_dt)
+
+        for substep in 1:nsubsteps
+            fast_time = slow_stage_time + (substep - 1) * fast_dt
+            dostep!(Qfast, fast, param, fast_time)
+            if split.coupled
+                cummulate_fast_solution!(
+                    slow_bl,
+                    fast_bl,
+                    fast.rhs!,
+                    Qfast,
+                    fast_time,
+                    fast_dt,
+                    substep,
+                )
+            end
+        end
+
+        # reconcile slow equation using fast equation
+        if split.coupled
+            reconcile_from_fast_to_slow!(
+                slow_bl,
+                fast_bl,
+                slow.rhs!,
+                fast.rhs!,
+                Qslow,
+                Qfast,
+            )
+        end
+    end
+    updatedt!(fast, fast_dt_in)
+
+    return nothing
+end

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

@@ -11,8 +11,11 @@ using ..MPIStateArrays
 using ..Mesh.Filters: apply!
 using ..Mesh.Grids: VerticalDirection
 using ..BalanceLaws: BalanceLaw
-import ..BalanceLaws:
-    nodal_update_auxiliary_state!,
+import ..BalanceLaws: nodal_update_auxiliary_state!
+using ..DGMethods.NumericalFluxes: RusanovNumericalFlux
+
+import ..DGMethods.NumericalFluxes: update_penalty!
+import ..DGMethods:
     vars_state_conservative,
     vars_state_auxiliary,
     vars_state_gradient,

src/Ocean/SplitExplicit/SplitExplicitModel.jl

@@ -0,0 +1,36 @@
+module SplitExplicit
+
+using ..HydrostaticBoussinesq
+using ..ShallowWater
+
+import ..BalanceLaws:
+    initialize_states!,
+    tendency_from_slow_to_fast!,
+    cummulate_fast_solution!,
+    reconcile_from_fast_to_slow!
+
+@inline initialize_states!(
+    slow::HydrostaticBoussinesqModel,
+    fast::ShallowWaterModel,
+    _...,
+) = nothing
+
+@inline tendency_from_slow_to_fast!(
+    slow::HydrostaticBoussinesqModel,
+    fast::ShallowWaterModel,
+    _...,
+) = nothing
+
+@inline cummulate_fast_solution!(
+    slow::HydrostaticBoussinesqModel,
+    fast::ShallowWaterModel,
+    _...,
+) = nothing
+
+@inline reconcile_from_fast_to_slow!(
+    slow::HydrostaticBoussinesqModel,
+    fast::ShallowWaterModel,
+    _...,
+) = nothing
+
+end