diff --git a/.gitignore b/.gitignore index e84627101..3cc979bbb 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,6 @@ Manifest.toml # no model output folders run_*/ + +# no video outputs +*.mp4 diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl index 74777bc97..14dd3dff8 100644 --- a/src/SpeedyWeather.jl +++ b/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") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 5f7dea0c6..432ee309d 100644 --- a/src/abstract_types.jl +++ b/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 diff --git a/src/dynamics/diagnostic_variables.jl b/src/dynamics/diagnostic_variables.jl index 141ba6484..a256c4d64 100644 --- a/src/dynamics/diagnostic_variables.jl +++ b/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 """ diff --git a/src/dynamics/models.jl b/src/dynamics/models.jl index 63b143998..d0d86c998 100644 --- a/src/dynamics/models.jl +++ b/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 diff --git a/src/dynamics/tendencies_dynamics.jl b/src/dynamics/tendencies_dynamics.jl index 9a38c0d72..3c4e815da 100644 --- a/src/dynamics/tendencies_dynamics.jl +++ b/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 diff --git a/src/dynamics/vertical_advection.jl b/src/dynamics/vertical_advection.jl new file mode 100644 index 000000000..a8ee7c767 --- /dev/null +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 4ddbc9acf..36fcbb080 100644 --- a/test/runtests.jl +++ b/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") diff --git a/test/vertical_advection.jl b/test/vertical_advection.jl new file mode 100644 index 000000000..b00f699f2 --- /dev/null +++ b/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 \ No newline at end of file