Try #1257:

slurmci-test.toml

@@ -40,7 +40,8 @@ cpu_gpu = [
   { file = "test/Numerics/SystemSolvers/bandedsystem.jl", slurmargs = ["--ntasks=3"], args = [] },
   { file = "test/Numerics/Mesh/interpolation.jl", slurmargs = ["--ntasks=3"], args = [] },
   { file = "test/Ocean/ShallowWater/GyreDriver.jl", slurmargs = ["--ntasks=1"], args = [] },
-  { file = "experiments/OceanBoxGCM/hydrostatic_spindown.jl", slurmargs = ["--ntasks=3"], args = [] },
+  { file = "experiments/OceanBoxGCM/hydrostatic_spindown.jl", slurmargs = ["--ntasks=1"], args = [] },
+  { file = "testSplitExplicit/hydrostatic_spindown.jl", slurmargs = ["--ntasks=1"], args = [] },
   { file = "tutorials/Numerics/DGMethods/nonnegative.jl", slurmargs = ["--ntasks=3"], args = [] },
 ]
 

src/ClimateMachine.jl

@@ -30,6 +30,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,28 @@ struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD}
     diffusion_direction::DD
     modeldata::MD
 end
+
+function basic_grid_info(dg::DGModel)
+    grid = dg.grid
+    topology = grid.topology
+
+    dim = dimensionality(grid)
+    N = polynomialorder(grid)
+
+    Nq = N + 1
+    Nqk = dim == 2 ? 1 : Nq
+    Nfp = Nq * Nqk
+    Np = dofs_per_element(grid)
+
+    nelem = length(topology.elems)
+    nvertelem = topology.stacksize
+    nhorzelem = div(nelem, nvertelem)
+    nrealelem = length(topology.realelems)
+    nhorzrealelem = div(nrealelem, nvertelem)
+
+    return (Nq, Nqk, Nfp, Np, nvertelem, nhorzelem, nhorzrealelem, nrealelem)
+end
+
 function DGModel(
     balance_law,
     grid,

src/Numerics/DGMethods/balance_law_interface.jl

@@ -305,3 +305,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/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,192 @@
+export SplitExplicitSolver
+
+using ClimateMachine.DGMethods:
+    initialize_states!,
+    tendency_from_slow_to_fast!,
+    cummulate_fast_solution!,
+    reconcile_from_fast_to_slow!
+
+LSRK2N = LowStorageRungeKutta2N
+
+"""
+    SplitExplicitMethod(slow_solver, fast_solver; dt, t0 = 0)
+
+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.
+
+### References
+"""
+mutable struct SplitExplicitSolver{SS, FS, RT} <: AbstractODESolver
+    "slow solver"
+    slow_solver::SS
+    "fast solver"
+    fast_solver::FS
+    "time step"
+    dt::RT
+    "time"
+    t::RT
+    "param to couple the models"
+    coupled::Bool
+    "param to step past slow dt"
+    add_fast_substeps::RT
+
+    function SplitExplicitSolver(
+        slow_solver::LSRK2N,
+        fast_solver,
+        Q = nothing;
+        dt = getdt(slow_solver),
+        t0 = slow_solver.t,
+        coupled = true,
+        add_fast_substeps = 0,
+    ) where {AT <: AbstractArray}
+        SS = typeof(slow_solver)
+        FS = typeof(fast_solver)
+        RT = real(eltype(slow_solver.dQ))
+        return new{SS, FS, RT}(
+            slow_solver,
+            fast_solver,
+            RT(dt),
+            RT(t0),
+            coupled,
+            RT(add_fast_substeps),
+        )
+    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 = similar(dQslow)
+
+    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 tentency adjustment before evalution of slow mode
+        if split.coupled
+            initialize_states!(
+                slow_bl,
+                fast_bl,
+                slow.rhs!,
+                fast.rhs!,
+                Qslow,
+                Qfast,
+            )
+        end
+
+        # 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
+        if split.coupled
+            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)
+            #  ---> need to cumulate U at this time (and not at each RKB sub-time-step)
+            if split.coupled
+                cummulate_fast_solution!(
+                    slow_bl,
+                    fast_bl,
+                    fast.rhs!,
+                    Qfast,
+                    fast_time,
+                    fast_dt,
+                    substep,
+                    # fast_steps,
+                    # fast_time_rec,
+                )
+            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,
+                # fast_time_rec,
+            )
+        end
+    end
+    updatedt!(fast, fast_dt_in)
+
+    return nothing
+end

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

@@ -22,7 +22,6 @@ using ..DGMethods.NumericalFluxes: RusanovNumericalFlux
 
 import ..DGMethods.NumericalFluxes: update_penalty!
 import ..DGMethods:
-
     vars_state_conservative,
     init_state_conservative!,
     vars_state_auxiliary,

src/Ocean/ShallowWater/ShallowWaterModel.jl

@@ -1,6 +1,6 @@
 module ShallowWater
 
-export SWModel, SWProblem
+export ShallowWaterModel, ShallowWaterProblem
 
 using StaticArrays
 using ..VariableTemplates
@@ -31,7 +31,7 @@ using ..DGMethods.NumericalFluxes
 ⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)
 ⊗(a::SVector, b::SVector) = a * b'
 
-abstract type SWProblem end
+abstract type ShallowWaterProblem end
 
 abstract type TurbulenceClosure end
 struct LinearDrag{L} <: TurbulenceClosure
@@ -44,14 +44,16 @@ end
 abstract type AdvectionTerm end
 struct NonLinearAdvection <: AdvectionTerm end
 
-struct SWModel{PS, P, T, A, S} <: BalanceLaw
+struct ShallowWaterModel{PS, P, T, A, S} <: BalanceLaw
     param_set::PS
     problem::P
     turbulence::T
     advection::A
     c::S
 end
 
+SWModel = ShallowWaterModel
+
 function vars_state_conservative(m::SWModel, T)
     @vars begin
         U::SVector{3, T}
@@ -68,7 +70,7 @@ end
 
 function vars_state_gradient(m::SWModel, T)
     @vars begin
-        U::SVector{3, T}
+        ∇U::SVector{3, T}
     end
 end
 
@@ -85,14 +87,12 @@ end
     α::Vars,
     t::Real,
 )
-    FT = eltype(q)
-    _grav::FT = grav(m.param_set)
     U = q.U
     η = q.η
     H = m.problem.H
 
     F.η += U
-    F.U += _grav * H * η * I
+    F.U += grav(m.param_set) * H * η * I
 
     advective_flux!(m, m.advection, F, q, α, t)
 
@@ -236,7 +236,7 @@ end
 
 function shallow_init_state! end
 function init_state_conservative!(m::SWModel, state::Vars, aux::Vars, coords, t)
-    shallow_init_state!(m.problem, m.turbulence, state, aux, coords, t)
+    shallow_init_state!(m, m.problem, state, aux, coords, t)
 end
 
 function boundary_state!(

src/Ocean/SplitExplicit/SplitExplicitModel.jl

@@ -0,0 +1,36 @@
+module SplitExplicit
+
+using ..HydrostaticBoussinesq
+using ..ShallowWater
+
+import ..DGMethods:
+    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