add auxiliary balance laws for integration

now getting caught on the dimension mismatch between SWModel and HBModel 
for velocity

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

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

@@ -14,9 +14,9 @@ using ..DGMethods:
     BalanceLaw,
     LocalGeometry,
     DGModel,
+    nodal_update_auxiliary_state!,
     indefinite_stack_integral!,
     reverse_indefinite_stack_integral!,
-    nodal_update_auxiliary_state!,
     copy_stack_field_down!
 using ..DGMethods.NumericalFluxes: RusanovNumericalFlux
 

src/Ocean/SplitExplicit/Communication.jl

@@ -0,0 +1,139 @@
+import ..DGMethods:
+    initialize_states!,
+    tendency_from_slow_to_fast!,
+    cummulate_fast_solution!,
+    reconcile_from_fast_to_slow!
+
+@inline function initialize_states!(
+    baroclinic::HydrostaticBoussinesqModel,
+    barotropic::ShallowWaterModel,
+    model_bc,
+    model_bt,
+    state_bc,
+    state_bt,
+)
+    model_bt.state_auxiliary.ηᶜ .= -0
+    model_bt.state_auxiliary.Uᶜ .= (@SVector [-0, -0, -0])'
+
+    state_bt.η .= model_bt.state_auxiliary.ηˢ
+    state_bt.U .= model_bt.state_auxiliary.Uˢ
+
+    model_bc.state_auxiliary.ΔGᵘ .= -0
+
+    return nothing
+end
+
+@inline function tendency_from_slow_to_fast!(
+    baroclinic::HydrostaticBoussinesqModel,
+    barotropic::ShallowWaterModel,
+    model_bc,
+    model_bt,
+    state_bc,
+    state_bt,
+    forcing_tendency,
+)
+    # integrate the tendency
+    tendency_model = model_bc.modeldata.tendency_model
+    update_auxiliary_state!(
+        tendency_model,
+        tendency_model.balance_law,
+        forcing_tendency,
+        0,
+    )
+
+    Nq, Nqk, _, _, nelemv, _, nelemh, _ = basic_grid_info(model_bc)
+
+    ## get top value (=integral over full depth) of ∫du
+    ∫du = tendency_model.state_auxiliary.∫du
+    boxy_∫du = reshape(∫du, Nq^2, Nqk, 2, nelemv, nelemh)
+    flat_∫du = @view boxy_∫du[:, end, :, end, :]
+
+    ## copy into Gᵁ of dgFast
+    model_bt.state_auxiliary.Gᵁ .= reshape(flat_∫du, Nq^2, 2, nelemh)
+
+    ## scale by -1/H and copy back to ΔGu
+    # note: since tendency_dg.auxstate.∫du is not used after this, could be
+    # re-used to store a 3-D copy of "-Gu"
+    ΔGᵘ = model_bc.state_auxiliary.ΔGᵘ
+    boxy_ΔGᵘ = reshape(ΔGᵘ, Nq^2, Nqk, 2, nelemv, nelemh)
+    boxy_ΔGᵘ .= -reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H
+
+    return nothing
+end
+
+@inline function cummulate_fast_solution!(
+    baroclinic::HydrostaticBoussinesqModel,
+    barotropic::ShallowWaterModel,
+    model_bt,
+    state_bt,
+    fast_time,
+    fast_dt,
+    substep,
+)
+    ### for now no averaging, just save the final time step
+    model_bt.state_auxiliary.ηᶜ .= state_bt.η
+    model_bt.state_auxiliary.Uᶜ .= state_bt.U
+
+    ### technically unnecessary right now, but for future changes
+    model_bt.state_auxiliary.ηˢ .= state_bt.η
+    model_bt.state_auxiliary.Uˢ .= state_bt.U
+
+    return nothing
+end
+
+@inline function reconcile_from_fast_to_slow!(
+    baroclinic::HydrostaticBoussinesqModel,
+    barotropic::ShallowWaterModel,
+    model_bc,
+    model_bt,
+    state_bc,
+    state_bt,
+)
+    Nq, Nqk, _, _, nelemv, _, nelemh, _ = basic_grid_info(model_bc)
+
+    # need to calculate int_u using integral kernels
+    # u_slow := u_slow + (1/H) * (u_fast - \int_{-H}^{0} u_slow)
+
+    # Compute: \int_{-H}^{0} u_slow)
+    ### need to make sure this is stored into aux.∫u
+
+    # integrate vertically horizontal velocity
+    flow_model = model_bc.modeldata.flow_model
+    update_auxiliary_state!(flow_model, flow_model.balance_law, state_bc, 0)
+
+    ## get top value (=integral over full depth)
+    ∫u = flow_model.state_auxiliary.∫u
+    boxy_∫u = reshape(∫u, Nq^2, Nqk, 2, nelemv, nelemh)
+    flat_∫u = @view boxy_∫u[:, end, :, end, :]
+
+    ## get time weighted averaged out of cumulative arrays
+    ## no weights right now
+    state_bt.state_auxiliary.Uᶜ .*= 1 # / fast_time_rec[2]
+    state_bt.state_auxiliary.ηᶜ .*= 1 # / fast_time_rec[2]
+
+    ### substract ∫u from U and divide by H
+
+    ### Δu is a place holder for 1/H * (Ū - ∫u)
+    Δu = model_bt.state_auxiliary.Δu
+    Δu .= 1 / baroclinic.problem.H * (model_bt.state_auxiliary.Uᶜ - flat_∫u)
+
+    ### copy the 2D contribution down the 3D solution
+    ### need to reshape these things for the broadcast
+    boxy_u = reshape(state_bc.u, Nq^2, Nqk, 2, nelemv, nelemh)
+    boxy_Δu = reshape(Δu, Nq^2, 1, 3, 1, nelemh)
+
+    ### copy 2D eta over to 3D model
+    boxy_η_3D = reshape(state_bc.η, Nq^2, Nqk, nelemv, nelemh)
+    boxy_η_2D = reshape(model_bt.state_auxiliary.ηᶜ, Nq^2, 1, 1, nelemh)
+
+    ### this works, we tested it
+    # currently SWModel has a 3-vector for U
+    # need to shrink to 2-vector at some point
+
+    #### turn off broadcast for testing purposes
+    #### probably need to add a flag here
+    # boxy_u .+= boxy_Δu[:, :, 1:2, :, :]
+    # boxy_η_3D .= boxy_η_2D
+
+    return nothing
+end

src/Ocean/SplitExplicit/FlowIntegralModel.jl

@@ -0,0 +1,62 @@
+struct FlowIntegralModel{M} <: BalanceLaw
+    ocean::M
+    function FlowIntegralModel(ocean::M) where {M}
+        return new{M}(ocean)
+    end
+end
+
+vars_state_gradient(tm::FlowIntegralModel, FT) = @vars()
+vars_state_gradient_flux(tm::FlowIntegralModel, FT) = @vars()
+
+vars_state_conservative(tm::FlowIntegralModel, FT) =
+    vars_state_conservative(tm.ocean, FT)
+
+function vars_state_auxiliary(m::FlowIntegralModel, T)
+    @vars begin
+        ∫u::SVector{2, T}
+    end
+end
+
+init_state_auxiliary!(tm::FlowIntegralModel, A::Vars, geom::LocalGeometry) =
+    nothing
+
+function vars_integrals(m::FlowIntegralModel, T)
+    @vars begin
+        ∫u::SVector{2, T}
+    end
+end
+
+@inline function integral_load_auxiliary_state!(
+    m::FlowIntegralModel,
+    I::Vars,
+    Q::Vars,
+    A::Vars,
+)
+    I.∫u = Q.u
+
+    return nothing
+end
+
+@inline function integral_set_auxiliary_state!(
+    m::FlowIntegralModel,
+    A::Vars,
+    I::Vars,
+)
+    A.∫u = I.∫u
+
+    return nothing
+end
+
+function update_auxiliary_state!(
+    dg::DGModel,
+    fm::FlowIntegralModel,
+    Q::MPIStateArray,
+    t::Real,
+)
+    A = dg.state_auxiliary
+
+    # compute vertical integral of u
+    indefinite_stack_integral!(dg, fm, Q, A, t) # bottom -> top
+
+    return true
+end

src/Ocean/SplitExplicit/SplitExplicitModel.jl

@@ -1,144 +1,33 @@
 module SplitExplicit
 
+using StaticArrays
+
 using ..HydrostaticBoussinesq
 using ..ShallowWater
-
-using StaticArrays
+using ..VariableTemplates
+using ..MPIStateArrays
+using ..DGMethods:
+    BalanceLaw,
+    LocalGeometry,
+    DGModel,
+    nodal_update_auxiliary_state!,
+    indefinite_stack_integral!,
+    basic_grid_info
 
 import ..DGMethods:
-    initialize_states!,
-    tendency_from_slow_to_fast!,
-    cummulate_fast_solution!,
-    reconcile_from_fast_to_slow!
-
-@inline function initialize_states!(
-    baroclinic::HydrostaticBoussinesqModel,
-    barotropic::ShallowWaterModel,
-    model_bc,
-    model_bt,
-    state_bc,
-    state_bt,
-)
-    model_bt.state_auxiliary.ηᶜ .= -0
-    model_bt.state_auxiliary.Uᶜ .= (@SVector [-0, -0, -0])'
-
-    state_bt.η .= model_bt.state_auxiliary.ηˢ
-    state_bt.U .= model_bt.state_auxiliary.Uˢ
-
-    model_bc.state_auxiliary.ΔGᵘ .= -0
-
-    return nothing
-end
-
-@inline function tendency_from_slow_to_fast!(
-    baroclinic::HydrostaticBoussinesqModel,
-    barotropic::ShallowWaterModel,
-    model_bc,
-    model_bt,
-    state_bc,
-    state_bt,
-    forcing_tendency,
-)
-    # integrate the tendency
-    tendency_model = model_bc.modeldata.tendency_model
-    update_auxiliary_state!(
-        tendency_model,
-        tendency_model.balance_law,
-        forcing_tendency,
-        0,
-    )
-
-    Nq, Nqk, _, _, nelemv, _, nelemh, _ = basic_grid_info(model_bc)
-
-    ## get top value (=integral over full depth) of ∫du
-    ∫du = tendency_model.state_auxiliary.∫du
-    boxy_∫du = reshape(∫du, Nq^2, Nqk, 2, nelemv, nelemh)
-    flat_∫du = @view boxy_∫du[:, end, :, end, :]
-
-    ## copy into Gᵁ of dgFast
-    model_bt.state_auxiliary.Gᵁ .= reshape(flat_∫du, Nq^2, 2, nelemh)
-
-    ## scale by -1/H and copy back to ΔGu
-    # note: since tendency_dg.auxstate.∫du is not used after this, could be
-    # re-used to store a 3-D copy of "-Gu"
-    ΔGᵘ = model_bc.state_auxiliary.ΔGᵘ
-    boxy_ΔGᵘ = reshape(ΔGᵘ, Nq^2, Nqk, 2, nelemv, nelemh)
-    boxy_ΔGᵘ .= -reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H
-
-    return nothing
-end
-
-@inline function cummulate_fast_solution!(
-    baroclinic::HydrostaticBoussinesqModel,
-    barotropic::ShallowWaterModel,
-    model_bt,
-    state_bt,
-    fast_time,
-    fast_dt,
-    substep,
-)
-    ### for now no averaging, just save the final time step
-    model_bt.state_auxiliary.ηᶜ .= state_bt.η
-    model_bt.state_auxiliary.Uᶜ .= state_bt.U
-
-    ### technically unnecessary right now, but for future changes
-    model_bt.state_auxiliary.ηˢ .= state_bt.η
-    model_bt.state_auxiliary.Uˢ .= state_bt.U
-
-    return nothing
-end
-
-@inline function reconcile_from_fast_to_slow!(
-    baroclinic::HydrostaticBoussinesqModel,
-    barotropic::ShallowWaterModel,
-    model_bc,
-    model_bt,
-    state_bc,
-    state_bt,
-)
-    Nq, Nqk, _, _, nelemv, _, nelemh, _ = basic_grid_info(model_bc)
-
-    # need to calculate int_u using integral kernels
-    # u_slow := u_slow + (1/H) * (u_fast - \int_{-H}^{0} u_slow)
-
-    # Compute: \int_{-H}^{0} u_slow)
-    ### need to make sure this is stored into aux.∫u
-
-    # integrate vertically horizontal velocity
-    flow_model = model_bc.modeldata.flow_model
-    update_auxiliary_state!(flow_model, flow_model.balance_law, state_bc, 0)
-
-    ## get top value (=integral over full depth)
-    ∫u = flow_model.state_auxiliary.∫u
-    boxy_∫u = reshape(∫u, Nq^2, Nqk, 2, nelemv, nelemh)
-    flat_∫u = @view boxy_∫u[:, end, :, end, :]
-
-    ## get time weighted averaged out of cumulative arrays
-    ## no weights right now
-    state_bt.state_auxiliary.Uᶜ .*= 1 # / fast_time_rec[2]
-    state_bt.state_auxiliary.ηᶜ .*= 1 # / fast_time_rec[2]
-
-    ### substract ∫u from U and divide by H
-
-    ### Δu is a place holder for 1/H * (Ū - ∫u)
-    Δu = model_bt.state_auxiliary.Δu
-    Δu .= 1 / baroclinic.problem.H * (model_bt.state_auxiliary.Uᶜ - flat_∫u)
-
-    ### copy the 2D contribution down the 3D solution
-    ### need to reshape these things for the broadcast
-    boxy_u = reshape(state_bc.u, Nq^2, Nqk, 2, nelemv, nelemh)
-    boxy_Δu = reshape(Δu, Nq^2, 1, 3, 1, nelemh)
-    ### this works, we tested it
-    # currently SWModel has a 3-vector for U
-    # need to shrink to 2-vector at some point
-    boxy_u .+= boxy_Δu[:, :, 1:2, :, :]
-
-    ### copy 2D eta over to 3D model
-    boxy_η_3D = reshape(state_bc.η, Nq^2, Nqk, nelemv, nelemh)
-    boxy_η_2D = reshape(model_bt.state_auxiliary.ηᶜ, Nq^2, 1, 1, nelemh)
-    boxy_η_3D .= boxy_η_2D
-
-    return nothing
-end
+    vars_state_conservative,
+    init_state_conservative!,
+    vars_state_auxiliary,
+    init_state_auxiliary!,
+    vars_state_gradient,
+    vars_state_gradient_flux,
+    vars_integrals,
+    integral_load_auxiliary_state!,
+    integral_set_auxiliary_state!,
+    update_auxiliary_state!
+
+include("Communication.jl")
+include("TendencyIntegralModel.jl")
+include("FlowIntegralModel.jl")
 
 end