Different vertical advection schemes

.gitignore

@@ -30,3 +30,6 @@ Manifest.toml
 
 # no model output folders
 run_*/
+
+# no video outputs
+*.mp4

src/SpeedyWeather.jl

@@ -155,6 +155,7 @@ include("dynamics/forcing.jl")
 include("dynamics/geopotential.jl")
 include("dynamics/initial_conditions.jl")
 include("dynamics/horizontal_diffusion.jl")
+include("dynamics/vertical_advection.jl")
 include("dynamics/implicit.jl")
 include("dynamics/scaling.jl")
 include("dynamics/tendencies.jl")

src/abstract_types.jl

@@ -26,6 +26,9 @@ abstract type AbstractColumnVariables{NF} end
 # FORCING (Barotropic and ShallowWaterModel)
 abstract type AbstractForcing{NF} end
 
+# VERTICAL ADVECTION (PrimitiveEquation)
+abstract type VerticalAdvection{NF,B} end
+
 # PARAMETERIZATIONS
 abstract type AbstractParameterization{NF} end
 abstract type BoundaryLayerDrag{NF} <: AbstractParameterization{NF} end

src/dynamics/diagnostic_variables.jl

@@ -46,11 +46,14 @@ struct GridVariables{NF<:AbstractFloat,Grid<:AbstractGrid{NF}}
     vor_grid            ::Grid  # vorticity
     div_grid            ::Grid  # divergence
     temp_grid           ::Grid  # absolute temperature [K]
+    temp_grid_prev      ::Grid  # absolute temperature of previous time step [K]
     temp_virt_grid      ::Grid  # virtual tempereature [K]  
     humid_grid          ::Grid  # specific_humidity
     geopot_grid         ::Grid  # geopotential (is that needed?)
     u_grid              ::Grid  # zonal velocity *coslat [m/s]
     v_grid              ::Grid  # meridional velocity *coslat [m/s]
+    u_grid_prev         ::Grid  # zonal velocity *coslat of previous time step [m/s]
+    v_grid_prev         ::Grid  # meridional velocity *coslat of previous time step [m/s]
 end
 
 function Base.zeros(::Type{GridVariables},SG::SpectralGrid)
@@ -59,14 +62,17 @@ function Base.zeros(::Type{GridVariables},SG::SpectralGrid)
     vor_grid            = zeros(Grid{NF},nlat_half)   # vorticity
     div_grid            = zeros(Grid{NF},nlat_half)   # divergence
     temp_grid           = zeros(Grid{NF},nlat_half)   # absolute temperature
+    temp_grid_prev      = zeros(Grid{NF},nlat_half)   # absolute temperature
     temp_virt_grid      = zeros(Grid{NF},nlat_half)   # virtual temperature
     humid_grid          = zeros(Grid{NF},nlat_half)   # specific humidity
     geopot_grid         = zeros(Grid{NF},nlat_half)   # geopotential
     u_grid              = zeros(Grid{NF},nlat_half)   # zonal velocity *coslat
     v_grid              = zeros(Grid{NF},nlat_half)   # meridonal velocity *coslat
+    u_grid_prev         = zeros(Grid{NF},nlat_half)   # zonal velocity *coslat
+    v_grid_prev         = zeros(Grid{NF},nlat_half)   # meridonal velocity *coslat
 
-    return GridVariables(   vor_grid,div_grid,temp_grid,temp_virt_grid,humid_grid,geopot_grid,
-                            u_grid,v_grid)
+    return GridVariables(   vor_grid,div_grid,temp_grid,temp_grid_prev,temp_virt_grid,humid_grid,geopot_grid,
+                            u_grid,v_grid,u_grid_prev,v_grid_prev)
 end
 
 """

src/dynamics/models.jl

@@ -139,6 +139,7 @@ Base.@kwdef struct PrimitiveDryModel{NF<:AbstractFloat, D<:AbstractDevice} <: Pr
     boundary_layer_drag::BoundaryLayerDrag{NF} = LinearDrag(spectral_grid)
     temperature_relaxation::TemperatureRelaxation{NF} = HeldSuarez(spectral_grid)
     static_energy_diffusion::VerticalDiffusion{NF} = StaticEnergyDiffusion(spectral_grid)
+    vertical_advection::VerticalAdvection{NF} = CenteredVerticalAdvection(spectral_grid)
     # vertical_diffusion::VerticalDiffusion{NF} = VerticalLaplacian(spectral_grid)
 
     # NUMERICS
@@ -206,6 +207,8 @@ Base.@kwdef struct PrimitiveWetModel{NF<:AbstractFloat, D<:AbstractDevice} <: Pr
     temperature_relaxation::TemperatureRelaxation{NF} = HeldSuarez(spectral_grid)
     static_energy_diffusion::VerticalDiffusion{NF} = StaticEnergyDiffusion(spectral_grid)
     large_scale_condensation::AbstractCondensation{NF} = SpeedyCondensation(spectral_grid)
+    vertical_advection::VerticalAdvection{NF} = CenteredVerticalAdvection(spectral_grid)
+
     # vertical_diffusion::VerticalDiffusion{NF} = VerticalLaplacian(spectral_grid)
 
     # NUMERICS

src/dynamics/tendencies_dynamics.jl

@@ -162,83 +162,6 @@ function vertical_velocity!(
         (uv∇lnp_sum_above .+ Δσₖ*uv∇lnp)        # and level k σₖ-weighted uv∇lnp here
 end
 
-function vertical_advection!(   layer::DiagnosticVariablesLayer,
-                                diagn::DiagnosticVariables,
-                                model::PrimitiveEquation)
-
-    (; k ) = layer       # which layer are we on?
-
-    wet_core = model isa PrimitiveWet
-    (; σ_levels_thick, nlev ) = model.geometry
-
-    # for k==1 "above" term is 0, for k==nlev "below" term is zero
-    # avoid out-of-bounds indexing with k_above, k_below as follows
-    k_above = max(1,k-1)        # just saturate, because M_1/2 = 0 (which zeros that term)
-    k_below = min(k+1,nlev)     # just saturate, because M_nlev+1/2 = 0 (which zeros that term)
-
-    # mass fluxes, M_1/2 = M_nlev+1/2 = 0, but k=1/2 isn't explicitly stored
-    σ_tend_above = diagn.layers[k_above].dynamics_variables.σ_tend
-    σ_tend_below = layer.dynamics_variables.σ_tend
-
-    # layer thickness Δσ on level k
-    Δσₖ = σ_levels_thick[k]
-
-    u_tend = layer.tendencies.u_tend_grid                   # zonal velocity
-    u_above = diagn.layers[k_above].grid_variables.u_grid
-    u = layer.grid_variables.u_grid
-    u_below = diagn.layers[k_below].grid_variables.u_grid
-
-    _vertical_advection!(u_tend,σ_tend_above,σ_tend_below,u_above,u,u_below,Δσₖ)
-
-    v_tend = layer.tendencies.v_tend_grid                   # meridional velocity
-    v_above = diagn.layers[k_above].grid_variables.v_grid
-    v = layer.grid_variables.v_grid
-    v_below = diagn.layers[k_below].grid_variables.v_grid
-
-    _vertical_advection!(v_tend,σ_tend_above,σ_tend_below,v_above,v,v_below,Δσₖ)
-
-    T_tend = layer.tendencies.temp_tend_grid                   # temperature
-    T_above = diagn.layers[k_above].grid_variables.temp_grid
-    T = layer.grid_variables.temp_grid
-    T_below = diagn.layers[k_below].grid_variables.temp_grid
-
-    _vertical_advection!(T_tend,σ_tend_above,σ_tend_below,T_above,T,T_below,Δσₖ)
-
-    if wet_core
-        q_tend = layer.tendencies.humid_tend_grid               # humidity
-        q_above = diagn.layers[k_above].grid_variables.humid_grid
-        q = layer.grid_variables.humid_grid
-        q_below = diagn.layers[k_below].grid_variables.humid_grid
-
-        _vertical_advection!(q_tend,σ_tend_above,σ_tend_below,q_above,q,q_below,Δσₖ)
-    end
-end
-
-# MULTI THREADED VERSION only writes into layer k
-function _vertical_advection!(  ξ_tend::Grid,           # tendency of quantity ξ at k
-                                σ_tend_above::Grid,     # vertical velocity at k-1/2
-                                σ_tend_below::Grid,     # vertical velocity at k+1/2
-                                ξ_above::Grid,          # quantity ξ at k-1
-                                ξ::Grid,                # quantity ξ at k
-                                ξ_below::Grid,          # quantity ξ at k+1
-                                Δσₖ::NF                 # layer thickness on σ levels
-                                ) where {NF<:AbstractFloat,Grid<:AbstractGrid{NF}}
-    # Δσₖ⁻¹ = 1/Δσₖ                                      # precompute
-
-    # # += as the tendencies already contain the parameterizations
-    # for ij in eachgridpoint(ξ_tend,σ_tend_above,σ_tend_below,ξ_above,ξ,ξ_below)
-    #     # 1st order upwind scheme
-    #     σ̇⁺ = max(σ_tend_above[ij],0)        # velocity into layer k from above
-    #     σ̇⁻ = min(σ_tend_below[ij],0)        # velocity into layer k from below
-    #     ξ_tend[ij] -= Δσₖ⁻¹ * (σ̇⁺*(ξ[ij] - ξ_above[ij]) + σ̇⁻*(ξ_below[ij] - ξ[ij]))
-    # end
-
-    # centred scheme
-    Δσₖ2⁻¹ = -1/2Δσₖ                                    # precompute
-    # += as the tendencies already contain the parameterizations
-    @. ξ_tend += Δσₖ2⁻¹ * (σ_tend_below*(ξ_below - ξ) + σ_tend_above*(ξ - ξ_above))
-end
-
 """
 $(TYPEDSIGNATURES)
 Function barrier to unpack `model`."""
@@ -750,9 +673,14 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer,
 
     (; vor_grid, div_grid, u_grid, v_grid ) = diagn.grid_variables
     (; temp_grid, humid_grid ) = diagn.grid_variables
+    (; temp_grid_prev, u_grid_prev, v_grid_prev) = diagn.grid_variables
     U = diagn.dynamics_variables.a      # reuse work arrays for velocities spectral
     V = diagn.dynamics_variables.b      # U = u*coslat, V=v*coslat
 
+    @. temp_grid_prev = temp_grid
+    @. u_grid_prev = u_grid
+    @. v_grid_prev = v_grid
+
     S = model.spectral_transform
     wet_core = model isa PrimitiveWet
 

src/dynamics/vertical_advection.jl

@@ -0,0 +1,145 @@
+# Dispersive and diffusive advection schemes `NF` is the type, `B` the half-stencil size
+abstract type DiffusiveVerticalAdvection{NF, B}  <: VerticalAdvection{NF, B} end
+abstract type DispersiveVerticalAdvection{NF, B} <: VerticalAdvection{NF, B} end
+
+struct UpwindVerticalAdvection{NF, B}   <: DiffusiveVerticalAdvection{NF, B} end
+struct WENOVerticalAdvection{NF}        <: DiffusiveVerticalAdvection{NF, 3} end
+struct CenteredVerticalAdvection{NF, B} <: DispersiveVerticalAdvection{NF, B} end
+
+CenteredVerticalAdvection(spectral_grid; order = 2) = CenteredVerticalAdvection{spectral_grid.NF, order       ÷ 2}()
+UpwindVerticalAdvection(spectral_grid; order = 5)   =   UpwindVerticalAdvection{spectral_grid.NF, (order + 1) ÷ 2}()
+WENOVerticalAdvection(spectral_grid)                =     WENOVerticalAdvection{spectral_grid.NF}()
+
+@inline retrieve_previous_time_step(variables, var) = getproperty(variables, Symbol(var, :_grid_prev))
+@inline  retrieve_current_time_step(variables, var) = getproperty(variables, Symbol(var, :_grid))
+
+@inline retrieve_time_step(::DiffusiveVerticalAdvection,  variables, var) = retrieve_previous_time_step(variables, var)
+@inline retrieve_time_step(::DispersiveVerticalAdvection, variables, var) =  retrieve_current_time_step(variables, var)
+
+@inline function retrieve_current_stencil(k, layers, var, nlev, ::VerticalAdvection{NF, B}) where {NF, B}
+    k_stencil = max.(min.(nlev, k-B:k+B), 1)
+    ξ_stencil = Tuple(retrieve_current_time_step(layers[k].grid_variables, var) for k in k_stencil)
+    return ξ_stencil
+end
+
+@inline function retrieve_previous_stencil(k, layers, var, nlev, ::VerticalAdvection{NF, B}) where {NF, B}
+    k_stencil = max.(min.(nlev, k-B:k+B), 1)
+    ξ_stencil = Tuple(retrieve_current_time_step(layers[k].grid_variables, var) for k in k_stencil)
+    return ξ_stencil
+end
+
+@inline retrieve_stencil(k, layers, var, nlev, scheme::DiffusiveVerticalAdvection)  = retrieve_previous_stencil(k, layers, var, nlev, scheme)
+@inline retrieve_stencil(k, layers, var, nlev, scheme::DispersiveVerticalAdvection) =  retrieve_current_stencil(k, layers, var, nlev, scheme)
+
+function vertical_advection!(   layer::DiagnosticVariablesLayer,
+                                diagn::DiagnosticVariables,
+                                model::PrimitiveEquation)
+
+    (; k ) = layer       # which layer are we on?
+
+    wet_core = model isa PrimitiveWet
+    (; σ_levels_thick, nlev ) = model.geometry
+
+    scheme = model.vertical_advection
+
+    # for k==1 "above" term is 0, for k==nlev "below" term is zero
+    # avoid out-of-bounds indexing with k_above, k_below as follows
+    k⁻ = max(1,k-1) # just saturate, because M_1/2 = 0 (which zeros that term)
+
+    # mass fluxes, M_1/2 = M_nlev+1/2 = 0, but k=1/2 isn't explicitly stored
+    σ_tend_above = diagn.layers[k⁻].dynamics_variables.σ_tend
+    σ_tend_below = layer.dynamics_variables.σ_tend
+
+    # layer thickness Δσ on level k
+    Δσₖ = σ_levels_thick[k]
+
+    for var in (:u, :v, :temp)
+        ξ_tend = getproperty(layer.tendencies, Symbol(var, :_tend_grid))
+        ξ_sten = retrieve_stencil(k, diagn.layers, var, nlev, scheme)
+        ξ      = retrieve_time_step(scheme, layer.grid_variables, var)
+
+        _vertical_advection!(ξ_tend,σ_tend_above,σ_tend_below,ξ_sten,ξ,Δσₖ,scheme)
+    end
+
+    if wet_core
+        ξ_tend = getproperty(layer.tendencies, :humid_tend_grid)
+        ξ_sten = retrieve_current_stencil(k, diagn.layers, :humid, nlev, scheme)
+        ξ      = retrieve_current_time_step(layer.grid_variables, :humid)
+
+        _vertical_advection!(ξ_tend,σ_tend_above,σ_tend_below,ξ_sten,ξ,Δσₖ,scheme)
+    end
+end
+
+# MULTI THREADED VERSION only writes into layer k
+function _vertical_advection!(  ξ_tend::Grid,                  # tendency of quantity ξ at k
+                                σ_tend_above::Grid,            # vertical velocity at k-1/2
+                                σ_tend_below::Grid,            # vertical velocity at k+1/2
+                                ξ_sten,                        # ξ stencil for vertical advection (from k-B to k+B)
+                                ξ::Grid,                       # ξ at level k
+                                Δσₖ::NF,                       # layer thickness on σ levels
+                                adv::VerticalAdvection{NF, B}  # vertical advection scheme
+                                ) where {NF<:AbstractFloat,Grid<:AbstractGrid{NF}, B}
+    Δσₖ⁻¹ = 1/Δσₖ                                      # precompute
+
+    # += as the tendencies already contain the parameterizations
+    for ij in eachgridpoint(ξ_tend)
+        σ̇⁻ = σ_tend_above[ij]       # velocity into layer k from above
+        σ̇⁺ = σ_tend_below[ij]       # velocity out of layer k to below
+
+        ξᶠ⁺ = reconstructed_at_face(ij, adv, σ̇⁺, ξ_sten[2:end])
+        ξᶠ⁻ = reconstructed_at_face(ij, adv, σ̇⁻, ξ_sten[1:end-1])
+
+        ξ_tend[ij] -=  Δσₖ⁻¹ * (σ̇⁺ * ξᶠ⁺ - σ̇⁻ * ξᶠ⁻ - ξ[ij] * (σ̇⁺ - σ̇⁻))
+    end
+end
+
+@inline reconstructed_at_face(ij, ::UpwindVerticalAdvection{NF, 1}, u, ξ) where NF = ifelse(u > 0, ξ[1][ij], ξ[2][ij])
+@inline reconstructed_at_face(ij, ::UpwindVerticalAdvection{NF, 2}, u, ξ) where NF = ifelse(u > 0, (2ξ[1][ij] + 5ξ[2][ij] - ξ[3][ij]) / 6,
+                                                                                                   (2ξ[4][ij] + 5ξ[3][ij] - ξ[2][ij]) / 6)
+@inline reconstructed_at_face(ij, ::UpwindVerticalAdvection{NF, 3}, u, ξ) where NF = ifelse(u > 0, (2ξ[1][ij] - 13ξ[2][ij] + 47ξ[3][ij] + 27ξ[4][ij] - 3ξ[5][ij]) / 60,
+                                                                                                   (2ξ[6][ij] - 13ξ[5][ij] + 47ξ[4][ij] + 27ξ[3][ij] - 3ξ[2][ij]) / 60)
+
+@inline reconstructed_at_face(ij, ::CenteredVerticalAdvection{NF, 1}, u, ξ) where NF = ( ξ[1][ij] +  ξ[2][ij]) / 2
+@inline reconstructed_at_face(ij, ::CenteredVerticalAdvection{NF, 2}, u, ξ) where NF = (-ξ[1][ij] + 7ξ[2][ij] + 7ξ[3][ij] - ξ[4][ij]) / 12
+
+const ε  = 1e-6
+const d₀ = 3/10
+const d₁ = 3/5
+const d₂ = 1/10
+
+@inline weight_β₀(S, NF) = NF(13/12) * (S[1] - 2S[2] + S[3])^2 + NF(1/4) * (3S[1] - 4S[2] +  S[3])^2
+@inline weight_β₁(S, NF) = NF(13/12) * (S[1] - 2S[2] + S[3])^2 + NF(1/4) * ( S[1]         -  S[3])^2
+@inline weight_β₂(S, NF) = NF(13/12) * (S[1] - 2S[2] + S[3])^2 + NF(1/4) * ( S[1] - 4S[2] + 3S[3])^2
+
+@inline p₀(S) = (2S[1] + 5S[2] -   S[3]) / 6 # downind stencil
+@inline p₁(S) = (-S[1] + 5S[2] +  2S[3]) / 6 # upwind stencil
+@inline p₂(S) = (2S[1] - 7S[2] + 11S[3]) / 6 # extrapolating stencil
+
+@inline τ₅(β₀, β₁, β₂) = abs(β₂ - β₀)
+
+@inline function weno_reconstruction(S₀, S₁, S₂, NF)
+    β₀ = weight_β₀(S₀, NF)
+    β₁ = weight_β₁(S₁, NF)
+    β₂ = weight_β₂(S₂, NF)
+
+    w₀ = NF(d₀) * (1 + (τ₅(β₀, β₁, β₂) / (β₀ + NF(ε)))^2)
+    w₁ = NF(d₁) * (1 + (τ₅(β₀, β₁, β₂) / (β₁ + NF(ε)))^2)
+    w₂ = NF(d₂) * (1 + (τ₅(β₀, β₁, β₂) / (β₂ + NF(ε)))^2)
+
+    w₀, w₁, w₂ = (w₀, w₁, w₂) ./ (w₀ + w₁ + w₂) 
+
+    return p₀(S₀) * w₀ + p₁(S₁) * w₁ + p₂(S₂) * w₂
+end
+
+@inline function reconstructed_at_face(ij, ::WENOVerticalAdvection{NF}, u, ξ) where NF
+    if u > 0
+        S₀ = (ξ[3][ij], ξ[4][ij], ξ[5][ij])
+        S₁ = (ξ[2][ij], ξ[3][ij], ξ[4][ij])
+        S₂ = (ξ[1][ij], ξ[2][ij], ξ[3][ij])
+    else
+        S₀ = (ξ[4][ij], ξ[3][ij], ξ[2][ij])
+        S₁ = (ξ[5][ij], ξ[4][ij], ξ[3][ij])
+        S₂ = (ξ[6][ij], ξ[5][ij], ξ[4][ij])
+    end
+    return weno_reconstruction(S₀, S₁, S₂, NF)
+end

test/runtests.jl

@@ -18,6 +18,7 @@ include("spectral_gradients.jl")
 # DYNAMICS
 include("diffusion.jl")
 include("time_stepping.jl")
+include("vertical_advection.jl")
 
 # VERTICAL LEVELS
 include("vertical_levels.jl")

test/vertical_advection.jl

@@ -0,0 +1,15 @@
+using SpeedyWeather: WENOVerticalAdvection, CenteredVerticalAdvection, UpwindVerticalAdvection
+using SpeedyWeather: vertical_advection!
+
+@testset "Vertical advection runs" begin
+    spectral_grid = SpectralGrid()
+    model_types = (PrimitiveDryModel, PrimitiveWetModel)
+    advection_schems = (WENOVerticalAdvection, CenteredVerticalAdvection, UpwindVerticalAdvection)
+
+    for Model in model_types, VerticalAdvection in advection_schems
+        model = Model(; spectral_grid, vertical_advection = VerticalAdvection(spectral_grid))
+        simulation = initialize!(model)
+        run!(simulation,n_days=10)
+        @test simulation.model.feedback.nars_detected == false    
+    end
+end