Merge #1268

1268: Couple the HBModel and SWModel for use in the Split-Explicit Stepper r=blallen a=blallen

# Description

A continuation of #1257, adding the necessary communication between the two models needed to have a true split-explicit stepper. The work in this PR enables us to run with up to a 90-minute timestep for the hydrostatic spindown test case. 

The broadcasts of reshaped MPIStateArrays depends on #1071 to run on the GPU. I'll pull over the commits from that PR for running bors, but we need to merge that before we merge this one.

There's quite a bit of partial work in this PR that was necessary to get things running, but that either isn't fully implemented everywhere it should be or isn't implemented super elegantly (mainly regarding collecting all the Ocean models into a single module and how multiple dispatch is used for the couple). I'll revisit that once we have the full timestepping machinery running.



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

experiments/OceanBoxGCM/refvals/ocean_gyre_refvals.jl

@@ -182,7 +182,6 @@ varr = [
         1.98256782258218657716e-05,
     ],
     [
-        "s_aux",
         :pkin,
         -9.01203660410079376852e-01,
         0.00000000000000000000e+00,
@@ -196,17 +195,6 @@ varr = [
         3.89428424525183031852e-05,
         4.05688221636862481177e-10,
         1.10561449663397378925e-05,
-    ],
-]
-parr = [
-    ["Q", "u[1]", 12, 12, 12, 12],
-    ["Q", "u[2]", 12, 12, 12, 12],
-    ["Q", :η, 12, 12, 12, 12],
-    ["Q", :θ, 12, 12, 12, 12],
-    ["s_aux", :y, 12, 12, 12, 12],
-    ["s_aux", :w, 12, 12, 12, 12],
-    ["s_aux", :pkin, 12, 12, 12, 12],
-    ["s_aux", :wz0, 12, 12, 8, 12],
 ]
 # END SCPRINT
 # SC ====================================================================================

src/Numerics/ODESolvers/SplitExplicitMethod.jl

@@ -36,8 +36,6 @@ mutable struct SplitExplicitSolver{SS, FS, RT, MSA} <: AbstractODESolver
     dt::RT
     "time"
     t::RT
-    "flag to couple the models"
-    coupled
     "storage for transfer tendency"
     dQ2fast::MSA
 
@@ -47,20 +45,19 @@ mutable struct SplitExplicitSolver{SS, FS, RT, MSA} <: AbstractODESolver
         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)
+        dQ2fast .= -0
         MSA = typeof(dQ2fast)
+
         return new{SS, FS, RT, MSA}(
             slow_solver,
             fast_solver,
             RT(dt),
             RT(t0),
-            coupled,
             dQ2fast,
         )
     end
@@ -95,53 +92,28 @@ function dostep!(
 
         # 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
+        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,
+        )
 
         # 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]
@@ -159,30 +131,42 @@ function dostep!(
         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!(
+            cummulate_fast_solution!(
                 slow_bl,
                 fast_bl,
-                slow.rhs!,
                 fast.rhs!,
-                Qslow,
                 Qfast,
+                fast_time,
+                fast_dt,
+                substep,
             )
         end
+
+        # 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)
+
+        # reconcile slow equation using fast equation
+        reconcile_from_fast_to_slow!(
+            slow_bl,
+            fast_bl,
+            slow.rhs!,
+            fast.rhs!,
+            Qslow,
+            Qfast,
+        )
     end
     updatedt!(fast, fast_dt_in)