Add traits to modify HBModel and SWModel for coupling

drive shallow water model by HB model

add coupling traits, use to replace forcing traits in shallow water model

move coupling for SWModel to file in SplitExplicit

add coupling to HBModel, get fully coupled model to converge to analytic solution

Add hydrostatic spindown tests with actual multi-rate happening

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

@@ -14,6 +14,7 @@ using ...Mesh.Geometry
 using ...DGMethods: DGModel, copy_stack_field_down!
 using ...DGMethods.NumericalFluxes
 using ...BalanceLaws
+using ..Ocean
 
 import ...DGMethods.NumericalFluxes: update_penalty!
 import ...BalanceLaws:
@@ -40,6 +41,7 @@ import ...BalanceLaws:
     reverse_indefinite_stack_integral!,
     reverse_integral_load_auxiliary_state!,
     reverse_integral_set_auxiliary_state!
+import ..Ocean: coriolis_force!
 
 ×(a::SVector, b::SVector) = StaticArrays.cross(a, b)
 ⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)
@@ -70,9 +72,10 @@ fₒ = first coriolis parameter (constant term)
     HydrostaticBoussinesqModel(problem)
 
 """
-struct HydrostaticBoussinesqModel{PS, P, T} <: BalanceLaw
+struct HydrostaticBoussinesqModel{PS, P, C, T} <: BalanceLaw
     param_set::PS
     problem::P
+    coupling::C
     ρₒ::T
     cʰ::T
     cᶻ::T
@@ -86,6 +89,7 @@ struct HydrostaticBoussinesqModel{PS, P, T} <: BalanceLaw
     function HydrostaticBoussinesqModel{FT}(
         param_set::PS,
         problem;
+        coupling = NotCoupled(),
         ρₒ = FT(1000),  # kg / m^3
         cʰ = FT(0),     # m/s
         cᶻ = FT(0),     # m/s
@@ -97,9 +101,10 @@ struct HydrostaticBoussinesqModel{PS, P, T} <: BalanceLaw
         fₒ = FT(1e-4),  # Hz
         β = FT(1e-11), # Hz / m
     ) where {FT <: AbstractFloat, PS}
-        return new{PS, typeof(problem), FT}(
+        return new{PS, typeof(problem), typeof(coupling), FT}(
             param_set,
             problem,
+            coupling,
             ρₒ,
             cʰ,
             cᶻ,
@@ -189,6 +194,7 @@ these are just copies in our model
 function vars_state_gradient(m::HBModel, T)
     @vars begin
         ∇u::SVector{2, T}
+        ∇uᵈ::SVector{2, T}
         ∇θ::T
     end
 end
@@ -207,9 +213,23 @@ this computation is done pointwise at each nodal point
 - `t`: time, not used
 """
 @inline function compute_gradient_argument!(m::HBModel, G::Vars, Q::Vars, A, t)
-    G.∇u = Q.u
     G.∇θ = Q.θ
 
+    velocity_gradient_argument!(m, m.coupling, G, Q, A, t)
+
+    return nothing
+end
+
+@inline function velocity_gradient_argument!(
+    m::HBModel,
+    ::NotCoupled,
+    G,
+    Q,
+    A,
+    t,
+)
+    G.∇u = Q.u
+
     return nothing
 end
 
@@ -248,15 +268,29 @@ this computation is done pointwise at each nodal point
     A::Vars,
     t,
 )
-    ν = viscosity_tensor(m)
-    D.ν∇u = -ν * G.∇u
+    velocity_gradient_flux!(m, m.coupling, D, G, Q, A, t)
 
     κ = diffusivity_tensor(m, G.∇θ[3])
     D.κ∇θ = -κ * G.∇θ
 
     return nothing
 end
 
+@inline function velocity_gradient_flux!(
+    m::HBModel,
+    ::NotCoupled,
+    D,
+    G,
+    Q,
+    A,
+    t,
+)
+    ν = viscosity_tensor(m)
+    D.ν∇u = -ν * G.∇u
+
+    return nothing
+end
+
 """
     viscosity_tensor(::HBModel)
 
@@ -419,38 +453,44 @@ t -> time, not used
     t::Real,
     direction,
 )
-    FT = eltype(Q)
-    _grav::FT = grav(m.param_set)
     @inbounds begin
-        u = Q.u # Horizontal components of velocity
-        η = Q.η
-        θ = Q.θ
-        w = A.w   # vertical velocity
-        pkin = A.pkin
+        # ∇h • (g η)
+        hydrostatic_pressure!(m, m.coupling, F, Q, A, t)
 
-        v = @SVector [u[1], u[2], w]
+        # ∇h • (- ∫(αᵀ θ))
+        pkin = A.pkin
         Iʰ = @SMatrix [
             1 -0
             -0 1
             -0 -0
         ]
-
-        # ∇h • (g η)
-        F.u += _grav * η * Iʰ
-
-        # ∇h • (- ∫(αᵀ θ))
-        F.u += _grav * pkin * Iʰ
+        F.u += grav(m.param_set) * pkin * Iʰ
 
         # ∇h • (v ⊗ u)
         # F.u += v * u'
 
         # ∇ • (u θ)
+        θ = Q.θ
+        v = @SVector [Q.u[1], Q.u[2], A.w]
         F.θ += v * θ
     end
 
     return nothing
 end
 
+@inline function hydrostatic_pressure!(m::HBModel, ::NotCoupled, F, Q, A, t)
+    η = Q.η
+    Iʰ = @SMatrix [
+        1 -0
+        -0 1
+        -0 -0
+    ]
+
+    F.u += grav(m.param_set) * η * Iʰ
+
+    return nothing
+end
+
 """
     flux_second_order!(::HBModel)
 
@@ -509,30 +549,33 @@ end
     t::Real,
     direction,
 )
-    @inbounds begin
-        u, v = Q.u # Horizontal components of velocity
-        wz0 = A.wz0
+    wz0 = A.wz0
+    S.η += wz0
 
-        # f × u
-        f = coriolis_force(m, A.y)
-        S.u -= @SVector [-f * v, f * u]
+    coriolis_force!(m, m.coupling, S, Q, A, t)
 
-        S.η += wz0
-    end
+    return nothing
+end
+
+@inline function coriolis_force!(m::HBModel, ::NotCoupled, S, Q, A, t)
+    # f × u
+    f = coriolis_parameter(m, A.y)
+    u, v = Q.u # Horizontal components of velocity
+    S.u -= @SVector [-f * v, f * u]
 
     return nothing
 end
 
 """
-    coriolis_force(::HBModel)
+    coriolis_parameter(::HBModel)
 
 northern hemisphere coriolis
 
 # Arguments
 - `m`: model object to dispatch on and get coriolis parameters
 - `y`: y-coordinate in the box
 """
-@inline coriolis_force(m::HBModel, y) = m.fₒ + m.β * y
+@inline coriolis_parameter(m::HBModel, y) = m.fₒ + m.β * y
 
 """
     wavespeed(::HBModel)
@@ -592,9 +635,13 @@ function update_auxiliary_state!(
         apply!(Q, (:θ,), dg.grid, exp_filter, direction = VerticalDirection())
     end
 
+    compute_flow_deviation!(dg, m, m.coupling, Q, t)
+
     return true
 end
 
+@inline compute_flow_deviation!(dg, ::HBModel, ::NotCoupled, _...) = nothing
+
 """
     update_auxiliary_state_gradient!(::HBModel)
 

src/Ocean/Ocean.jl

@@ -1,5 +1,14 @@
 module Ocean
 
+export AbstractOceanCoupling, NotCoupled, PartiallyCoupled, FullyCoupled
+
+abstract type AbstractOceanCoupling end
+struct NotCoupled <: AbstractOceanCoupling end
+struct PartiallyCoupled <: AbstractOceanCoupling end
+struct FullyCoupled <: AbstractOceanCoupling end
+
+function coriolis_force! end
+
 include("HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl")
 include("ShallowWater/ShallowWaterModel.jl")
 include("SplitExplicit/SplitExplicitModel.jl")

src/Ocean/ShallowWater/ShallowWaterModel.jl

@@ -11,6 +11,7 @@ using ...Mesh.Geometry
 using ...DGMethods
 using ...DGMethods.NumericalFluxes
 using ...BalanceLaws
+using ..Ocean
 
 import ...DGMethods.NumericalFluxes: update_penalty!
 import ...BalanceLaws:
@@ -27,6 +28,7 @@ import ...BalanceLaws:
     source!,
     wavespeed,
     boundary_state!
+import ..Ocean: coriolis_force!
 
 ×(a::SVector, b::SVector) = StaticArrays.cross(a, b)
 ⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)
@@ -45,9 +47,10 @@ end
 abstract type AdvectionTerm end
 struct NonLinearAdvection <: AdvectionTerm end
 
-struct ShallowWaterModel{PS, P, T, A, S} <: BalanceLaw
+struct ShallowWaterModel{PS, P, T, A, C, S} <: BalanceLaw
     param_set::PS
     problem::P
+    coupling::C
     turbulence::T
     advection::A
     c::S
@@ -159,12 +162,13 @@ function compute_gradient_flux!(
     α::Vars,
     t::Real,
 )
-    compute_gradient_flux!(m.turbulence, σ, δ, q, α, t)
+    compute_gradient_flux!(m, m.turbulence, σ, δ, q, α, t)
 end
 
-compute_gradient_flux!(::LinearDrag, _...) = nothing
+compute_gradient_flux!(::SWModel, ::LinearDrag, _...) = nothing
 
 @inline function compute_gradient_flux!(
+    ::SWModel,
     T::ConstantViscosity,
     σ::Vars,
     δ::Grad,
@@ -189,13 +193,15 @@ function flux_second_order!(
     α::Vars,
     t::Real,
 )
-    flux_second_order!(m.turbulence, G, q, σ, α, t)
+    flux_second_order!(m, m.turbulence, m.coupling, G, q, σ, α, t)
 end
 
-flux_second_order!(::LinearDrag, _...) = nothing
+flux_second_order!(::SWModel, ::LinearDrag, _...) = nothing
 
 @inline function flux_second_order!(
+    ::SWModel,
     ::ConstantViscosity,
+    ::NotCoupled,
     G::Grad,
     q::Vars,
     σ::Vars,
@@ -213,18 +219,29 @@ end
     m::SWModel{P},
     S::Vars,
     q::Vars,
-    diffusive::Vars,
+    d::Vars,
     α::Vars,
     t::Real,
     direction,
 ) where {P}
-    S.U += α.τ
+    coriolis_force!(m, m.coupling, S, q, α, t)
+    forcing_term!(m, m.coupling, S, q, α, t)
+    linear_drag!(m.turbulence, S, q, α, t)
+
+    return nothing
+end
 
-    f = α.f
-    U, V = q.U
+@inline function coriolis_force!(::SWModel, ::NotCoupled, S, Q, A, t)
+    # f × u
+    f = A.f
+    U, V = Q.U
     S.U -= @SVector [-f * V, f * U]
 
-    linear_drag!(m.turbulence, S, q, α, t)
+    return nothing
+end
+
+@inline function forcing_term!(::SWModel, ::NotCoupled, S, Q, A, t)
+    S.U += A.τ
 
     return nothing
 end