Skip to content

Commit

Permalink
Merge pull request #362 from SpeedyWeather/ss/vertical_upwinding
Browse files Browse the repository at this point in the history
Different vertical advection schemes
  • Loading branch information
milankl authored Aug 28, 2023
2 parents 59a8133 + cc3dfce commit 4952295
Show file tree
Hide file tree
Showing 9 changed files with 184 additions and 79 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,6 @@ Manifest.toml

# no model output folders
run_*/

# no video outputs
*.mp4
1 change: 1 addition & 0 deletions src/SpeedyWeather.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
3 changes: 3 additions & 0 deletions src/abstract_types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 8 additions & 2 deletions src/dynamics/diagnostic_variables.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

"""
Expand Down
3 changes: 3 additions & 0 deletions src/dynamics/models.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
82 changes: 5 additions & 77 deletions src/dynamics/tendencies_dynamics.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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`."""
Expand Down Expand Up @@ -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

Expand Down
145 changes: 145 additions & 0 deletions src/dynamics/vertical_advection.jl
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
15 changes: 15 additions & 0 deletions test/vertical_advection.jl
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 4952295

Please sign in to comment.