-
Notifications
You must be signed in to change notification settings - Fork 29
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add forcing and drag extensions to tests
- Loading branch information
Showing
2 changed files
with
199 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,196 @@ | ||
@testset "Extending forcing and drag" begin | ||
Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag{NF} | ||
|
||
# DIMENSIONS from SpectralGrid | ||
"Spectral resolution as max degree of spherical harmonics" | ||
trunc::Int | ||
|
||
# OPTIONS | ||
"Relaxation time scale τ" | ||
time_scale::Second = Day(6) | ||
|
||
"Jet strength [m/s]" | ||
u₀::Float64 = 20 | ||
|
||
"latitude of Gaussian jet [˚N]" | ||
latitude::Float64 = 30 | ||
|
||
"Width of Gaussian jet [˚]" | ||
width::Float64 = 6 | ||
|
||
# TO BE INITIALISED | ||
"Relaxation back to reference vorticity" | ||
ζ₀::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) | ||
end | ||
|
||
function JetDrag(SG::SpectralGrid;kwargs...) | ||
return JetDrag{SG.NF}(;SG.trunc,kwargs...) | ||
end | ||
|
||
function SpeedyWeather.initialize!( drag::JetDrag, | ||
model::SpeedyWeather.ModelSetup) | ||
|
||
(;spectral_grid, geometry) = model | ||
(;Grid,NF,nlat_half) = spectral_grid | ||
u = zeros(Grid{NF},nlat_half) | ||
|
||
lat = geometry.latds | ||
|
||
for ij in eachindex(u) | ||
u[ij] = drag.u₀ * exp(-(lat[ij]-drag.latitude)^2/(2*drag.width^2)) | ||
end | ||
|
||
û = SpeedyTransforms.spectral(u,one_more_degree=true) | ||
v̂ = zero(û) | ||
SpeedyTransforms.curl!(drag.ζ₀,û,v̂,model.spectral_transform) | ||
return nothing | ||
end | ||
|
||
function SpeedyWeather.drag!( diagn::SpeedyWeather.DiagnosticVariablesLayer, | ||
progn::SpeedyWeather.PrognosticVariablesLayer, | ||
drag::JetDrag, | ||
time::DateTime, | ||
model::SpeedyWeather.ModelSetup) | ||
|
||
(;vor) = progn | ||
(;vor_tend) = diagn.tendencies | ||
(;ζ₀) = drag | ||
|
||
(;radius) = model.spectral_grid | ||
r = radius/drag.time_scale.value | ||
for lm in eachindex(vor,vor_tend,ζ₀) | ||
vor_tend[lm] -= r*(vor[lm] - ζ₀[lm]) | ||
end | ||
|
||
SpeedyTransforms.spectral_truncation!(vor_tend) # set lmax+1 to zero | ||
|
||
return nothing | ||
end | ||
|
||
Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF} | ||
|
||
# DIMENSIONS from SpectralGrid | ||
"Spectral resolution as max degree of spherical harmonics" | ||
trunc::Int | ||
|
||
"Number of latitude rings, used for latitudinal mask" | ||
nlat::Int | ||
|
||
|
||
# OPTIONS | ||
"Decorrelation time scale τ [days]" | ||
decorrelation_time::Second = Day(2) | ||
|
||
"Stirring strength A [1/s²]" | ||
strength::Float64 = 1e-11 | ||
|
||
"Stirring latitude [˚N]" | ||
latitude::Float64 = 45 | ||
|
||
"Stirring width [˚]" | ||
width::Float64 = 24 | ||
|
||
|
||
# TO BE INITIALISED | ||
"Stochastic stirring term S" | ||
S::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) | ||
|
||
"a = A*sqrt(1 - exp(-2dt/τ)), the noise factor times the stirring strength [1/s²]" | ||
a::Base.RefValue{NF} = Ref(zero(NF)) | ||
|
||
"b = exp(-dt/τ), the auto-regressive factor [1]" | ||
b::Base.RefValue{NF} = Ref(zero(NF)) | ||
|
||
"Latitudinal mask, confined to mid-latitude storm track by default [1]" | ||
lat_mask::Vector{NF} = zeros(NF,nlat) | ||
end | ||
|
||
function StochasticStirring(SG::SpectralGrid;kwargs...) | ||
(;trunc,Grid,nlat_half) = SG | ||
nlat = RingGrids.get_nlat(Grid,nlat_half) | ||
return StochasticStirring{SG.NF}(;trunc,nlat,kwargs...) | ||
end | ||
|
||
function SpeedyWeather.initialize!( forcing::StochasticStirring, | ||
model::SpeedyWeather.ModelSetup) | ||
|
||
# precompute forcing strength, scale with radius^2 as is the vorticity equation | ||
(;radius) = model.spectral_grid | ||
A = radius^2 * forcing.strength | ||
|
||
# precompute noise and auto-regressive factor, packed in RefValue for mutability | ||
dt = model.time_stepping.Δt_sec | ||
τ = forcing.decorrelation_time.value # in seconds | ||
forcing.a[] = A*sqrt(1 - exp(-2dt/τ)) | ||
forcing.b[] = exp(-dt/τ) | ||
|
||
# precompute the latitudinal mask | ||
(;Grid,nlat_half) = model.spectral_grid | ||
latd = RingGrids.get_latd(Grid,nlat_half) | ||
|
||
for j in eachindex(forcing.lat_mask) | ||
# Gaussian centred at forcing.latitude of width forcing.width | ||
forcing.lat_mask[j] = exp(-(forcing.latitude-latd[j])^2/forcing.width^2*2) | ||
end | ||
|
||
return nothing | ||
end | ||
|
||
function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer, | ||
progn::SpeedyWeather.PrognosticVariablesLayer, | ||
forcing::StochasticStirring, | ||
time::DateTime, | ||
model::SpeedyWeather.ModelSetup) | ||
SpeedyWeather.forcing!(diagn,forcing,model.spectral_transform) | ||
end | ||
|
||
function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer, | ||
forcing::StochasticStirring{NF}, | ||
spectral_transform::SpectralTransform) where NF | ||
|
||
# noise and auto-regressive factors | ||
a = forcing.a[] # = sqrt(1 - exp(-2dt/τ)) | ||
b = forcing.b[] # = exp(-dt/τ) | ||
|
||
(;S) = forcing | ||
@inbounds for lm in eachindex(S) | ||
# Barnes and Hartmann, 2011 Eq. 2 | ||
Qi = 2rand(Complex{NF}) - (1 + im) # ~ [-1,1] in complex | ||
S[lm] = a*Qi + b*S[lm] | ||
end | ||
|
||
# to grid-point space | ||
S_grid = diagn.dynamics_variables.a_grid | ||
SpeedyTransforms.gridded!(S_grid,S,spectral_transform) | ||
|
||
# mask everything but mid-latitudes | ||
RingGrids._scale_lat!(S_grid,forcing.lat_mask) | ||
|
||
# back to spectral space | ||
(;vor_tend) = diagn.tendencies | ||
SpeedyTransforms.spectral!(vor_tend,S_grid,spectral_transform) | ||
SpeedyTransforms.spectral_truncation!(vor_tend) # set lmax+1 to zero | ||
|
||
return nothing | ||
end | ||
|
||
spectral_grid = SpectralGrid(trunc=31,nlev=1) | ||
|
||
drag = JetDrag(spectral_grid,time_scale=Day(6)) | ||
forcing = StochasticStirring(spectral_grid) | ||
initial_conditions = StartFromRest() | ||
|
||
# with barotropic model | ||
model = BarotropicModel(;spectral_grid,initial_conditions,forcing,drag) | ||
simulation = initialize!(model) | ||
|
||
run!(simulation,period=Day(5)) | ||
@test simulation.model.feedback.nars_detected == false | ||
|
||
# with shallow water model | ||
model = ShallowWaterModel(;spectral_grid,initial_conditions,forcing,drag) | ||
simulation = initialize!(model) | ||
|
||
run!(simulation,period=Day(5)) | ||
@test simulation.model.feedback.nars_detected == false | ||
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters