Add communication functions between SWModel and HBModel

fill out communication functions and add necessary aux variables

don't statecheck aux while adding new aux variables

turn off restart values for now

remove time averaging over fast mode and add flag for reconciliation

add zero initializations

replace .=- with .-=

rebase fix for `test_vertical_integral_model.jl`

convert `tendency_from_slow_to_fast!` to new reshape&broadcast setup

Update `reconcile_from_fast_to_slow!` with new reshape&broadcast method

Update to new functions from CliMA#1280

Co-Authored-By: Jean-Michel Campin <jm-c@users.noreply.github.com>

src/Numerics/ODESolvers/SplitExplicitMethod.jl

@@ -54,7 +54,9 @@ mutable struct SplitExplicitSolver{SS, FS, RT, MSA} <: AbstractODESolver
         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,

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

@@ -164,6 +164,8 @@ function vars_state(m::HBModel, ::Auxiliary, T)
         w::T     # ∫(-∇⋅u)
         pkin::T  # ∫(-αᵀθ)
         wz0::T   # w at z=0
+        uᵈ::SVector{2, T}    # velocity deviation from vertical mean
+        ΔGᵘ::SVector{2, T}   # vertically averaged tendency
     end
 end
 

src/Ocean/HydrostaticBoussinesq/SimpleBoxProblem.jl

@@ -52,6 +52,9 @@ function ocean_init_aux!(m::HBModel, p::AbstractSimpleBoxProblem, A, geom)
     A.pkin = 0
     A.wz0 = 0
 
+    A.uᵈ = @SVector [-0, -0]
+    A.ΔGᵘ = @SVector [-0, -0]
+
     return nothing
 end
 

src/Ocean/ShallowWater/ShallowWaterModel.jl

@@ -63,6 +63,8 @@ function vars_state(m::SWModel, ::Auxiliary, T)
     @vars begin
         f::T
         τ::SVector{2, T}  # value includes τₒ, g, and ρ
+        Gᵁ::SVector{2, T} # integral of baroclinic tendency
+        Δu::SVector{2, T} # reconciliation Δu = 1/H * (Ū - ∫u)
     end
 end
 
@@ -237,7 +239,7 @@ end
 
 function shallow_init_aux! end
 function init_state_auxiliary!(m::SWModel, aux::Vars, geom::LocalGeometry)
-    shallow_init_aux!(m.problem, aux, geom)
+    shallow_init_aux!(m, m.problem, aux, geom)
 end
 
 function shallow_init_state! end

src/Ocean/SplitExplicit/Communication.jl

@@ -0,0 +1,144 @@
+@inline function initialize_states!(
+    baroclinic::HBModel{C},
+    barotropic::SWModel{C},
+    model_bc,
+    model_bt,
+    state_bc,
+    state_bt,
+) where {C <: Coupled}
+    model_bc.state_auxiliary.ΔGᵘ .= -0
+
+    return nothing
+end
+
+@inline function tendency_from_slow_to_fast!(
+    baroclinic::HBModel{C},
+    barotropic::SWModel{C},
+    model_bc,
+    model_bt,
+    state_bc,
+    state_bt,
+    forcing_tendency,
+) where {C <: Coupled}
+    FT = eltype(state_bc)
+    Nq, Nqk, _, _, nelemv, nelemh, nrealelemh, _ = basic_grid_info(model_bc)
+
+    #### integrate the tendency
+    model_int = model_bc.modeldata.integral_model
+    integral = model_int.balance_law
+    update_auxiliary_state!(model_int, integral, forcing_tendency, 0)
+
+    ### properly shape MPIStateArrays
+    num_aux_int = number_states(integral, Auxiliary(), FT)
+    data_int = model_int.state_auxiliary.data
+    data_int = reshape(data_int, Nq^2, Nqk, num_aux_int, nelemv, nelemh)
+
+    num_aux_bt = number_states(barotropic, Auxiliary(), FT)
+    data_bt = model_bt.state_auxiliary.data
+    data_bt = reshape(data_bt, Nq^2, num_aux_bt, nelemh)
+
+    num_aux_bc = number_states(baroclinic, Auxiliary(), FT)
+    data_bc = model_bc.state_auxiliary.data
+    data_bc = reshape(data_bc, Nq^2, Nqk, num_aux_bc, nelemv, nelemh)
+
+    ### get vars indices
+    index_∫du = varsindex(vars_state(integral, Auxiliary(), FT), :(∫x))
+    index_Gᵁ = varsindex(vars_state(barotropic, Auxiliary(), FT), :Gᵁ)
+    index_ΔGᵘ = varsindex(vars_state(baroclinic, Auxiliary(), FT), :ΔGᵘ)
+
+    ### get top value (=integral over full depth) of ∫du
+    ∫du = @view data_int[:, end, index_∫du, end, 1:nrealelemh]
+
+    ### copy into Gᵁ of barotropic model
+    Gᵁ = @view data_bt[:, index_Gᵁ, 1:nrealelemh]
+    Gᵁ .= ∫du
+
+    ### get top value (=integral over full depth) of ∫du
+    ∫du = @view data_int[:, end:end, index_∫du, end:end, 1:nrealelemh]
+
+    ### save vertically averaged tendency to remove from 3D tendency
+    ### need to reshape for the broadcast
+    ΔGᵘ = @view data_bc[:, :, index_ΔGᵘ, :, 1:nrealelemh]
+    ΔGᵘ .-= ∫du / baroclinic.problem.H
+
+    return nothing
+end
+
+@inline function cummulate_fast_solution!(
+    baroclinic::HBModel{C},
+    barotropic::SWModel{C},
+    model_bt,
+    state_bt,
+    fast_time,
+    fast_dt,
+    substep,
+) where {C <: Coupled}
+    return nothing
+end
+
+@inline function reconcile_from_fast_to_slow!(
+    baroclinic::HBModel{C},
+    barotropic::SWModel{C},
+    model_bc,
+    model_bt,
+    state_bc,
+    state_bt,
+) where {C <: Coupled}
+    FT = eltype(state_bc)
+    Nq, Nqk, _, _, nelemv, nelemh, nrealelemh, _ = basic_grid_info(model_bc)
+
+    ### integrate the horizontal velocity
+    model_int = model_bc.modeldata.integral_model
+    integral = model_int.balance_law
+    update_auxiliary_state!(model_int, integral, state_bc, 0)
+
+    ### properly shape MPIStateArrays
+    num_aux_int = number_states(integral, Auxiliary(), FT)
+    data_int = model_int.state_auxiliary.data
+    data_int = reshape(data_int, Nq^2, Nqk, num_aux_int, nelemv, nelemh)
+
+    num_aux_bt = number_states(barotropic, Auxiliary(), FT)
+    data_bt_aux = model_bt.state_auxiliary.data
+    data_bt_aux = reshape(data_bt_aux, Nq^2, num_aux_bt, nelemh)
+
+    num_state_bt = number_states(barotropic, Prognostic(), FT)
+    data_bt_state = state_bt.data
+    data_bt_state = reshape(data_bt_state, Nq^2, num_state_bt, nelemh)
+
+    num_state_bc = number_states(baroclinic, Prognostic(), FT)
+    data_bc_state = state_bc.data
+    data_bc_state =
+        reshape(data_bc_state, Nq^2, Nqk, num_state_bc, nelemv, nelemh)
+
+    ### get vars indices
+    index_∫u = varsindex(vars_state(integral, Auxiliary(), FT), :(∫x))
+    index_Δu = varsindex(vars_state(barotropic, Auxiliary(), FT), :Δu)
+    index_U = varsindex(vars_state(barotropic, Prognostic(), FT), :U)
+    index_u = varsindex(vars_state(baroclinic, Prognostic(), FT), :u)
+    index_η_3D = varsindex(vars_state(baroclinic, Prognostic(), FT), :η)
+    index_η_2D = varsindex(vars_state(barotropic, Prognostic(), FT), :η)
+
+    ### get top value (=integral over full depth)
+    ∫u = @view data_int[:, end, index_∫u, end, 1:nrealelemh]
+
+    ### Δu is a place holder for 1/H * (Ū - ∫u)
+    Δu = @view data_bt_aux[:, index_Δu, 1:nrealelemh]
+    U = @view data_bt_state[:, index_U, 1:nrealelemh]
+    Δu .= 1 / baroclinic.problem.H * (U - ∫u)
+
+    ### copy the 2D contribution down the 3D solution
+    ### need to reshape for the broadcast
+    data_bt_aux = reshape(data_bt_aux, Nq^2, 1, num_aux_bt, 1, nelemh)
+    Δu = @view data_bt_aux[:, :, index_Δu, :, 1:nrealelemh]
+    u = @view data_bc_state[:, :, index_u, :, 1:nrealelemh]
+    u .+= Δu
+
+    ### copy η from barotropic mode to baroclinic mode
+    ### need to reshape for the broadcast
+    data_bt_state = reshape(data_bt_state, Nq^2, 1, num_state_bt, 1, nelemh)
+    η_2D = @view data_bt_state[:, :, index_η_2D, :, 1:nrealelemh]
+    η_3D = @view data_bc_state[:, :, index_η_3D, :, 1:nrealelemh]
+    η_3D .= η_2D
+
+    return nothing
+end

src/Ocean/SplitExplicit/SplitExplicitModel.jl

@@ -11,36 +11,14 @@ using ...Mesh.Geometry
 using ...DGMethods
 using ...BalanceLaws
 
-include("VerticalIntegralModel.jl")
-
 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
+include("VerticalIntegralModel.jl")
+include("Communication.jl")
 
 end

test/Ocean/ShallowWater/2D_hydrostatic_spindown.jl

@@ -62,7 +62,7 @@ function shallow_init_state!(
     return nothing
 end
 
-function shallow_init_aux!(p::SimpleBox, A, geom)
+function shallow_init_aux!(m::ShallowWaterModel, p::SimpleBox, A, geom)
     @inbounds y = geom.coord[2]
 
     FT = eltype(A)
@@ -202,7 +202,7 @@ function make_callbacks(
         if s
             starttime[] = now()
         else
-            energy = norm(Q_slow)
+            energy = norm(Q_fast)
             @info @sprintf(
                 """Update
                 simtime = %8.2f / %8.2f

test/Ocean/ShallowWater/GyreDriver.jl

@@ -76,7 +76,7 @@ function shallow_init_state!(
     end
 end
 
-function shallow_init_aux!(p::GyreInABox, aux, geom)
+function shallow_init_aux!(m::ShallowWaterModel, p::GyreInABox, aux, geom)
     @inbounds y = geom.coord[2]
 
     Lʸ = p.Lʸ

test/Ocean/SplitExplicit/hydrostatic_spindown.jl

@@ -67,16 +67,20 @@ function shallow_init_state!(
     return nothing
 end
 
-function shallow_init_aux!(p::GyreInABox, aux, geom)
+function shallow_init_aux!(m::ShallowWaterModel, p::GyreInABox, A, geom)
+    @inbounds x = geom.coord[1]
     @inbounds y = geom.coord[2]
 
     Lʸ = p.Lʸ
     τₒ = p.τₒ
     fₒ = p.fₒ
     β = p.β
 
-    aux.τ = @SVector [-τₒ * cos(π * y / Lʸ), 0]
-    aux.f = fₒ + β * (y - Lʸ / 2)
+    A.τ = @SVector [-τₒ * cos(π * y / Lʸ), 0]
+    A.f = fₒ + β * (y - Lʸ / 2)
+
+    A.Gᵁ = @SVector [-0, -0]
+    A.Δu = @SVector [-0, -0]
 
     return nothing
 end
@@ -325,7 +329,7 @@ vtkpath = "vtk_split"
 const timeend = 24 * 3600 # s
 const tout = 2 * 3600 # s
 # const timeend = 1200 # s
-# const tout = 600 # s
+# const tout = 300 # s
 
 const N = 4
 const Nˣ = 5

test/Ocean/SplitExplicit/test_vertical_integral_model.jl

@@ -67,7 +67,7 @@ function shallow_init_state!(
     return nothing
 end
 
-function shallow_init_aux!(p::GyreInABox, aux, geom)
+function shallow_init_aux!(m::ShallowWaterModel, p::GyreInABox, aux, geom)
     @inbounds y = geom.coord[2]
 
     Lʸ = p.Lʸ