From b122e2f059c57f7cae75ece746a7cbf70e419325 Mon Sep 17 00:00:00 2001 From: Maximilian Gelbrecht Date: Sat, 14 Dec 2024 01:17:35 +0100 Subject: [PATCH 1/2] RingGrinds more general indexing with Colon (less escapes) (#637) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * RingGrinds more general indexing with Colon (less escapes) * changelog RingGrids colon indexing --------- Co-authored-by: Milan Klöwer --- CHANGELOG.md | 1 + src/RingGrids/general.jl | 2 +- test/grids.jl | 3 +++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f113bbda..dad761940 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- RingGrids indexing with leading Colon should now always return another RingGrid instance [#637](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/637) - Roll back GPUArrays upgrade to ensure CUDA compatibility [#636](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/636) - Change default timestep to 40min at T31 [#623](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/623) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index cbfc52f4f..fedfb1b8b 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -76,7 +76,7 @@ Base.@propagate_inbounds Base.getindex(G::AbstractGridArray, ijk...) = getindex( @inline function Base.getindex( G::GridArray, col::Colon, - k::Integer..., + k..., ) where {GridArray<:AbstractGridArray} GridArray_ = nonparametric_type(GridArray) # obtain parameters from G.data return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings) diff --git a/test/grids.jl b/test/grids.jl index a11683b39..bbb1bd936 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -380,6 +380,9 @@ end @test grid[1:2, 1:2, 1:2] == grid.data[1:2, 1:2, 1:2] @test grid[1, 1, :] == grid.data[1, 1, :] + @test SpeedyWeather.RingGrids.nonparametric_type(typeof(grid[:,1:2,1:2])) <: RingGrids.nonparametric_type(G) + @test grid[:, 1:2, 1:2].data == grid.data[:, 1:2, 1:2] + idx = CartesianIndex((1, 2, 3)) @test grid[idx] == grid.data[idx] From f38f9184cd06d82d838bba1c9d5b9dadf4de5483 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Milan=20Kl=C3=B6wer?= Date: Sat, 14 Dec 2024 09:55:03 +0000 Subject: [PATCH 2/2] Forcing/drag for primitive models too (#635) * forcing/drag for primitive models too * tests adapted * update changelog * export StochasticStirring * include @testset * StartFromRest defined for PrimitiveEquation model * increase stochastic stirring strength * ConstantPressure no op for ShallowWater * docs: function signature updated --- CHANGELOG.md | 1 + docs/src/forcing_drag.md | 2 +- src/dynamics/drag.jl | 81 ++++++++++- src/dynamics/forcing.jl | 101 +++++++++++++- src/dynamics/initial_conditions.jl | 38 +++-- src/dynamics/random_process.jl | 8 +- src/dynamics/tendencies.jl | 12 +- src/dynamics/time_integration.jl | 2 + src/models/primitive_dry.jl | 6 + src/models/primitive_wet.jl | 6 + test/extending.jl | 216 ----------------------------- test/forcing_drag.jl | 34 +++++ test/runtests.jl | 2 +- 13 files changed, 261 insertions(+), 248 deletions(-) delete mode 100644 test/extending.jl create mode 100644 test/forcing_drag.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index dad761940..9e2931dc4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Forcing/drag for primitive models [#635](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/635) - RingGrids indexing with leading Colon should now always return another RingGrid instance [#637](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/637) - Roll back GPUArrays upgrade to ensure CUDA compatibility [#636](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/636) - Change default timestep to 40min at T31 [#623](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/623) diff --git a/docs/src/forcing_drag.md b/docs/src/forcing_drag.md index f7a69ca05..53cd45e21 100644 --- a/docs/src/forcing_drag.md +++ b/docs/src/forcing_drag.md @@ -197,8 +197,8 @@ function SpeedyWeather.forcing!( diagn::DiagnosticVariables, progn::PrognosticVariables, forcing::StochasticStirring, - model::AbstractModel, lf::Integer, + model::AbstractModel, ) # function barrier only forcing!(diagn, forcing, model.spectral_transform) diff --git a/src/dynamics/drag.jl b/src/dynamics/drag.jl index 71d4696a8..f5e05145c 100644 --- a/src/dynamics/drag.jl +++ b/src/dynamics/drag.jl @@ -1,5 +1,15 @@ abstract type AbstractDrag <: AbstractModelComponent end +# function barrier for all drags to unpack model.drag +function drag!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, + model::AbstractModel, +) + drag!(diagn, progn, model.drag, lf, model) +end + ## NO DRAG export NoDrag struct NoDrag <: AbstractDrag end @@ -9,8 +19,7 @@ initialize!(::NoDrag, ::AbstractModel) = nothing function drag!( diagn::DiagnosticVariables, progn::PrognosticVariables, drag::NoDrag, - model::AbstractModel, - lf::Integer) + args...) return nothing end @@ -37,8 +46,8 @@ function drag!( diagn::DiagnosticVariables, progn::PrognosticVariables, drag::QuadraticDrag, - model::AbstractModel, lf::Integer, + model::AbstractModel, ) drag!(diagn, drag) end @@ -71,4 +80,70 @@ function drag!( Fu[ij, k] -= c*speed*u[ij, k] # -= as the tendencies already contain forcing Fv[ij, k] -= c*speed*v[ij, k] end +end + +export JetDrag +@kwdef struct JetDrag{NF, SpectralVariable2D} <: SpeedyWeather.AbstractDrag + + # DIMENSIONS from SpectralGrid + "Spectral resolution as max degree of spherical harmonics" + trunc::Int + + "[OPTION] Relaxation time scale τ" + time_scale::Second = Day(6) + + "[OPTION] Jet strength [m/s]" + u₀::NF = 20 + + "[OPTION] latitude of Gaussian jet [˚N]" + latitude::NF = 30 + + "[OPTION] Width of Gaussian jet [˚]" + width::NF = 6 + + # TO BE INITIALISED + "Relaxation back to reference vorticity" + ζ₀::SpectralVariable2D = zeros(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1) +end + +function JetDrag(SG::SpectralGrid; kwargs...) + return JetDrag{SG.NF, SG.SpectralVariable2D}(; SG.trunc, kwargs...) +end + +function initialize!(drag::JetDrag, model::AbstractModel) + (; 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 + + û = transform(u, model.spectral_transform) + v̂ = zero(û) + curl!(drag.ζ₀, û, v̂, model.spectral_transform) + return nothing +end + +function drag!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + drag::JetDrag, + lf::Integer, + model::AbstractModel, +) + vor = progn.vor[lf] + (; vor_tend) = diagn.tendencies + (; ζ₀) = drag + + # scale by radius as is vorticity + (; radius) = model.spectral_grid + r = radius/drag.time_scale.value + + k = diagn.nlayers # drag only on surface layer + for lm in eachharmonic(vor_tend) + vor_tend[lm, k] -= r*(vor[lm, k] - ζ₀[lm]) + end end \ No newline at end of file diff --git a/src/dynamics/forcing.jl b/src/dynamics/forcing.jl index aa24b57ab..5e8731f8a 100644 --- a/src/dynamics/forcing.jl +++ b/src/dynamics/forcing.jl @@ -1,16 +1,27 @@ abstract type AbstractForcing <: AbstractModelComponent end +# function barrier for all forcings to unpack model.forcing +function forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + lf::Integer, + model::AbstractModel, +) + forcing!(diagn, progn, model.forcing, lf, model) +end + ## NO FORCING = dummy forcing export NoForcing struct NoForcing <: AbstractForcing end NoForcing(SG::SpectralGrid) = NoForcing() initialize!(::NoForcing, ::AbstractModel) = nothing -function forcing!( diagn::DiagnosticVariables, - progn::PrognosticVariables, - forcing::NoForcing, - model::AbstractModel, - lf::Integer) +function forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + forcing::NoForcing, + args..., +) return nothing end @@ -98,8 +109,8 @@ function forcing!( diagn::DiagnosticVariables, progn::PrognosticVariables, forcing::JetStreamForcing, - model::AbstractModel, lf::Integer, + model::AbstractModel, ) forcing!(diagn, forcing) end @@ -119,8 +130,84 @@ function forcing!( for (j, ring) in enumerate(eachring(Fu)) F = amplitude[j] for ij in ring - Fu[ij] = tapering[k]*F + # += to accumulate, not overwrite previous parameterizations/terms + Fu[ij] += tapering[k]*F end end end +end + +export StochasticStirring +@kwdef struct StochasticStirring{NF, VectorType} <: AbstractForcing + + "Number of latitude rings, used for latitudinal mask" + nlat::Int + + "[OPTION] Stirring strength A [1/s²]" + strength::NF = 1e-9 + + "[OPTION] Stirring latitude [˚N]" + latitude::NF = 45 + + "[OPTION] Stirring width [˚]" + width::NF = 24 + + # TO BE INITIALISED + "Latitudinal mask, confined to mid-latitude storm track by default [1]" + lat_mask::VectorType = zeros(NF, nlat) +end + +function StochasticStirring(SG::SpectralGrid; kwargs...) + return StochasticStirring{SG.NF, SG.VectorType}(; nlat=SG.nlat, kwargs...) +end + +function initialize!( + forcing::StochasticStirring, + model::AbstractModel) + + model.random_process isa NoRandomProcess && + @warn "StochasticStirring needs a random process. model.random_process is a NoRandomProcess." + + # precompute the latitudinal mask + (; latd) = model.geometry + 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 +end + +function forcing!( + diagn::DiagnosticVariables, + progn::PrognosticVariables, + forcing::StochasticStirring, + lf::Integer, + model::AbstractModel, +) + forcing!(diagn, forcing, model.spectral_transform) +end + +function forcing!( + diagn::DiagnosticVariables, + forcing::StochasticStirring, + spectral_transform::SpectralTransform, +) + # get random values from random process + S_grid = diagn.grid.random_pattern + + # mask everything but mid-latitudes + RingGrids._scale_lat!(S_grid, forcing.lat_mask) + + # back to spectral space + S_masked = diagn.dynamics.a_2D + transform!(S_masked, S_grid, spectral_transform) + + # scale by radius^2 as is the vorticity equation, and scale to forcing strength + S_masked .*= (diagn.scale[]^2 * forcing.strength) + + # force every layer + (; vor_tend) = diagn.tendencies + + for k in eachmatrix(vor_tend) + vor_tend[:, k] .+= S_masked + end end \ No newline at end of file diff --git a/src/dynamics/initial_conditions.jl b/src/dynamics/initial_conditions.jl index d3dea4063..e7133d43f 100644 --- a/src/dynamics/initial_conditions.jl +++ b/src/dynamics/initial_conditions.jl @@ -1,7 +1,7 @@ abstract type AbstractInitialConditions <: AbstractModelComponent end export InitialConditions -@kwdef struct InitialConditions{V,P,T,H} <: AbstractInitialConditions +@kwdef struct InitialConditions{V, P, T, H} <: AbstractInitialConditions vordiv::V = ZeroInitially() pres::P = ZeroInitially() temp::T = ZeroInitially() @@ -19,13 +19,13 @@ function initialize!( has(model, :humid) && initialize!(progn, IC.humid, model) end -InitialConditions(::Type{<:Barotropic}) = InitialConditions(;vordiv = StartWithRandomVorticity()) -InitialConditions(::Type{<:ShallowWater}) = InitialConditions(;vordiv = ZonalJet()) +InitialConditions(::Type{<:Barotropic}) = InitialConditions(; vordiv = StartWithRandomVorticity()) +InitialConditions(::Type{<:ShallowWater}) = InitialConditions(; vordiv = ZonalJet()) function InitialConditions(::Type{<:PrimitiveDry}) vordiv = ZonalWind() pres = PressureOnOrography() temp = JablonowskiTemperature() - return InitialConditions(;vordiv,pres,temp) + return InitialConditions(;vordiv, pres, temp) end function InitialConditions(::Type{<:PrimitiveWet}) @@ -33,7 +33,7 @@ function InitialConditions(::Type{<:PrimitiveWet}) pres = PressureOnOrography() temp = JablonowskiTemperature() humid = ConstantRelativeHumidity() - return InitialConditions(;vordiv,pres,temp,humid) + return InitialConditions(;vordiv, pres, temp, humid) end export ZeroInitially @@ -42,8 +42,23 @@ initialize!(::PrognosticVariables,::ZeroInitially,::AbstractModel) = nothing # to avoid a breaking change, like ZeroInitially export StartFromRest -struct StartFromRest <: AbstractInitialConditions end -initialize!(::PrognosticVariables,::StartFromRest,::AbstractModel) = nothing +@kwdef struct StartFromRest{P, T, H} <: AbstractInitialConditions + pres::P = ConstantPressure() + temp::T = JablonowskiTemperature() + humid::H = ZeroInitially() +end + +initialize!(::PrognosticVariables, ::StartFromRest, ::Barotropic) = nothing + +function initialize!( + progn::PrognosticVariables, + IC::StartFromRest, + model::AbstractModel, +) + has(model, :pres) && initialize!(progn, IC.pres, model) + has(model, :temp) && initialize!(progn, IC.temp, model) + has(model, :humid) && initialize!(progn, IC.humid, model) +end export StartWithRandomVorticity @@ -62,7 +77,7 @@ $(TYPEDSIGNATURES) Start with random vorticity as initial conditions""" function initialize!( progn::PrognosticVariables{NF}, initial_conditions::StartWithRandomVorticity, - model::AbstractModel) where NF + model::Barotropic) where NF lmax = progn.trunc + 1 power = initial_conditions.power + 1 # +1 as power is summed of orders m @@ -376,7 +391,7 @@ $(TYPEDSIGNATURES) Initial conditions from Jablonowski and Williamson, 2006, QJR Meteorol. Soc""" function initialize!( progn::PrognosticVariables{NF}, initial_conditions::JablonowskiTemperature, - model::AbstractModel) where NF + model::PrimitiveEquation) where NF (;u₀, η₀, ΔT, Tmin) = initial_conditions (;σ_tropopause) = initial_conditions @@ -554,6 +569,9 @@ function initialize!( progn::PrognosticVariables, return nothing end +# for shallow water constant pressure = 0 as pres=interface displacement here +initialize!(::PrognosticVariables, ::ConstantPressure, ::ShallowWater) = nothing + export ConstantRelativeHumidity @kwdef struct ConstantRelativeHumidity <: AbstractInitialConditions relhumid_ref::Float64 = 0.7 @@ -562,7 +580,7 @@ end function initialize!( progn::PrognosticVariables, IC::ConstantRelativeHumidity, - model::AbstractModel, + model::PrimitiveEquation, ) (; relhumid_ref) = IC (; nlayers, σ_levels_full) = model.geometry diff --git a/src/dynamics/random_process.jl b/src/dynamics/random_process.jl index 904c0366a..a8af10884 100644 --- a/src/dynamics/random_process.jl +++ b/src/dynamics/random_process.jl @@ -49,7 +49,7 @@ independently. Transformed after every time step to grid space with a `clamp` applied to limit extrema. For reproducability `seed` can be provided and an independent `random_number_generator` is used that is reseeded on every `initialize!`. Fields are $(TYPEDFIELDS)""" -@kwdef struct SpectralAR1Process{NF} <: AbstractRandomProcess +@kwdef struct SpectralAR1Process{NF, VectorType} <: AbstractRandomProcess trunc::Int "[OPTION] Time scale of the AR1 process" @@ -73,12 +73,12 @@ that is reseeded on every `initialize!`. Fields are $(TYPEDFIELDS)""" "Precomputed auto-regressive factor [1], function of time scale and model time step" autoregressive_factor::Base.RefValue{NF} = Ref(zero(NF)) - "Precomputed noise factors [1] for evert total wavenumber l" - noise_factors::Vector{NF} = zeros(NF, trunc+2) + "Precomputed noise factors [1] for every total wavenumber l" + noise_factors::VectorType = zeros(NF, trunc+2) end # generator function -SpectralAR1Process(SG::SpectralGrid; kwargs...) = SpectralAR1Process{SG.NF}(trunc=SG.trunc; kwargs...) +SpectralAR1Process(SG::SpectralGrid; kwargs...) = SpectralAR1Process{SG.NF, SG.VectorType}(trunc=SG.trunc; kwargs...) function initialize!( process::SpectralAR1Process, diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl index dff5443b2..67e22c9ec 100644 --- a/src/dynamics/tendencies.jl +++ b/src/dynamics/tendencies.jl @@ -6,9 +6,9 @@ function dynamics_tendencies!( lf::Integer, # leapfrog index to evaluate tendencies at model::Barotropic, ) - forcing!(diagn, progn, model.forcing, model, lf) # = (Fᵤ, Fᵥ) forcing for u, v - drag!(diagn, progn, model.drag, model, lf) # drag term for u, v - vorticity_flux!(diagn, model) # = ∇×(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ) + forcing!(diagn, progn, lf, model) # = (Fᵤ, Fᵥ) forcing for u, v + drag!(diagn, progn, lf, model) # drag term for u, v + vorticity_flux!(diagn, model) # = ∇×(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ) end """ @@ -20,12 +20,12 @@ function dynamics_tendencies!( lf::Integer, # leapfrog index to evaluate tendencies at model::ShallowWater, ) - (; forcing, drag, planet, atmosphere, orography) = model + (; planet, atmosphere, orography) = model (; spectral_transform, geometry) = model # for compatibility with other AbstractModels pressure pres = interface displacement η here - forcing!(diagn, progn, forcing, model, lf) # = (Fᵤ, Fᵥ, Fₙ) forcing for u, v, η - drag!(diagn, progn, drag, model, lf) # drag term for u, v + forcing!(diagn, progn, lf, model) # = (Fᵤ, Fᵥ, Fₙ) forcing for u, v, η + drag!(diagn, progn, lf, model) # drag term for u, v # = ∇×(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ), tendency for vorticity # = ∇⋅(v(ζ+f) + Fᵤ, -u(ζ+f) + Fᵥ), tendency for divergence diff --git a/src/dynamics/time_integration.jl b/src/dynamics/time_integration.jl index 074a83fe7..e5496a2df 100644 --- a/src/dynamics/time_integration.jl +++ b/src/dynamics/time_integration.jl @@ -314,6 +314,8 @@ function timestep!( end if model.dynamics # switch on/off all dynamics + forcing!(diagn, progn, lf2, model) + drag!(diagn, progn, lf2, model) dynamics_tendencies!(diagn, progn, lf2, model) # dynamical core implicit_correction!(diagn, model.implicit, progn) # semi-implicit time stepping corrections else # just transform physics tendencies to spectral space diff --git a/src/models/primitive_dry.jl b/src/models/primitive_dry.jl index 5c919e866..8ea9df8b5 100644 --- a/src/models/primitive_dry.jl +++ b/src/models/primitive_dry.jl @@ -22,6 +22,8 @@ $(TYPEDFIELDS)""" PA, # <:AbstractParticleAdvection, IC, # <:AbstractInitialConditions, RP, # <:AbstractRandomProcess, + FR, # <:AbstractForcing, + DR, # <:AbstractDrag, LS, # <:AbstractLandSeaMask, OC, # <:AbstractOcean, LA, # <:AbstractLand, @@ -61,6 +63,8 @@ $(TYPEDFIELDS)""" particle_advection::PA = NoParticleAdvection() initial_conditions::IC = InitialConditions(PrimitiveDry) random_process::RP = NoRandomProcess() + forcing::FR = NoForcing() + drag::DR = NoDrag() # BOUNDARY CONDITIONS orography::OR = EarthOrography(spectral_grid) @@ -117,6 +121,8 @@ function initialize!(model::PrimitiveDry; time::DateTime = DEFAULT_DATE) initialize!(model.geopotential, model) initialize!(model.adiabatic_conversion, model) initialize!(model.random_process, model) + initialize!(model.forcing, model) + initialize!(model.drag, model) # boundary conditionss initialize!(model.orography, model) diff --git a/src/models/primitive_wet.jl b/src/models/primitive_wet.jl index 8d0e1e17f..8be4c68b1 100644 --- a/src/models/primitive_wet.jl +++ b/src/models/primitive_wet.jl @@ -22,6 +22,8 @@ $(TYPEDFIELDS)""" PA, # <:AbstractParticleAdvection, IC, # <:AbstractInitialConditions, RP, # <:AbstractRandomProcess, + FR, # <:AbstractForcing, + DR, # <:AbstractDrag, LS, # <:AbstractLandSeaMask, OC, # <:AbstractOcean, LA, # <:AbstractLand, @@ -67,6 +69,8 @@ $(TYPEDFIELDS)""" particle_advection::PA = NoParticleAdvection() initial_conditions::IC = InitialConditions(PrimitiveWet) random_process::RP = NoRandomProcess() + forcing::FR = NoForcing() + drag::DR = NoDrag() # BOUNDARY CONDITIONS orography::OR = EarthOrography(spectral_grid) @@ -129,6 +133,8 @@ function initialize!(model::PrimitiveWet; time::DateTime = DEFAULT_DATE) initialize!(model.geopotential, model) initialize!(model.adiabatic_conversion, model) initialize!(model.random_process, model) + initialize!(model.forcing, model) + initialize!(model.drag, model) # boundary conditions initialize!(model.orography, model) diff --git a/test/extending.jl b/test/extending.jl deleted file mode 100644 index d4916397b..000000000 --- a/test/extending.jl +++ /dev/null @@ -1,216 +0,0 @@ -@testset "Extending forcing and drag" begin - @kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag - - # 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₀::NF = 20 - - "latitude of Gaussian jet [˚N]" - latitude::NF = 30 - - "Width of Gaussian jet [˚]" - width::NF = 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::AbstractModel) - - (; 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.transform(u, model.spectral_transform) - v̂ = zero(û) - SpeedyTransforms.curl!(drag.ζ₀, û, v̂, model.spectral_transform) - return nothing - end - - function SpeedyWeather.drag!( - diagn::DiagnosticVariables, - progn::PrognosticVariables, - drag::JetDrag, - model::AbstractModel, - lf::Integer, - ) - - vor = progn.vor[lf] - (; vor_tend) = diagn.tendencies - (; ζ₀) = drag - - (; radius) = model.spectral_grid - r = radius/drag.time_scale.value - - k = diagn.nlayers # drag only on surface layer - for lm in eachharmonic(vor_tend) - vor_tend[lm, k] -= r*(vor[lm, k] - ζ₀[lm]) - end - end - - @kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing - - # 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::NF = 1e-11 - - "Stirring latitude [˚N]" - latitude::NF = 45 - - "Stirring width [˚]" - width::NF = 24 - - "Minimum degree of spherical harmonics to force" - lmin::Int = 8 - - "Maximum degree of spherical harmonics to force" - lmax::Int = 40 - - "Minimum order of spherical harmonics to force" - mmin::Int = 4 - - "Maximum order of spherical harmonics to force" - mmax::Int = lmax - - # 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, nlat) = SG - return StochasticStirring{SG.NF}(; trunc, nlat, kwargs...) - end - - function SpeedyWeather.initialize!( forcing::StochasticStirring, - model::AbstractModel) - - # 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::DiagnosticVariables, - progn::PrognosticVariables, - forcing::StochasticStirring, - model::AbstractModel, - lf::Integer, - ) - SpeedyWeather.forcing!(diagn, forcing, model.spectral_transform) - end - - function SpeedyWeather.forcing!( - diagn::DiagnosticVariables, - 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 - lmax, mmax = size(S, as=Matrix) - - @inbounds for m in 1:mmax - for l in m:lmax - if (forcing.mmin <= m <= forcing.mmax) && - (forcing.lmin <= l <= forcing.lmax) - # Barnes and Hartmann, 2011 Eq. 2 - Qi = 2rand(Complex{NF}) - (1 + im) # ~ [-1, 1] in complex - S[l, m] = a*Qi + b*S[l, m] - end - end - end - - # to grid-point space - S_grid = diagn.dynamics.a_2D_grid - transform!(S_grid, S, spectral_transform) - - # mask everything but mid-latitudes - RingGrids._scale_lat!(S_grid, forcing.lat_mask) - - # back to spectral space - S_masked = diagn.dynamics.a_2D - transform!(S_masked, S_grid, spectral_transform) - k = diagn.nlayers # only force surface layer - diagn.tendencies.vor_tend[:, k] .+= S_masked - return nothing - end - - spectral_grid = SpectralGrid(trunc=31, nlayers=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 \ No newline at end of file diff --git a/test/forcing_drag.jl b/test/forcing_drag.jl new file mode 100644 index 000000000..ef9a90230 --- /dev/null +++ b/test/forcing_drag.jl @@ -0,0 +1,34 @@ +@testset "Forcing and drag: 2D" begin + # 2D models + spectral_grid = SpectralGrid(trunc=31, nlayers=1) + drag = JetDrag(spectral_grid, time_scale=Day(6)) + forcing = StochasticStirring(spectral_grid) + random_process = SpectralAR1Process(spectral_grid) + initial_conditions = StartFromRest() + + @testset for Model in (BarotropicModel, ShallowWaterModel) + model = Model(spectral_grid; initial_conditions, forcing, drag, random_process) + simulation = initialize!(model) + + run!(simulation, period=Day(15), output=true) + @test simulation.model.feedback.nars_detected == false + end +end + +@testset "Forcing and drag: 3D" begin + + # 3D models + spectral_grid = SpectralGrid(trunc=31, nlayers=8) + drag = JetDrag(spectral_grid, time_scale=Day(6)) + forcing = StochasticStirring(spectral_grid) + random_process = SpectralAR1Process(spectral_grid) + initial_conditions = StartFromRest() + + @testset for Model in (PrimitiveDryModel, PrimitiveWetModel) + model = Model(spectral_grid; initial_conditions, forcing, drag, random_process) + simulation = initialize!(model) + + run!(simulation, period=Day(15), output=true) + @test simulation.model.feedback.nars_detected == false + end +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index b590bc4e0..76f6e7925 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -31,6 +31,7 @@ include("time_stepping.jl") include("vertical_advection.jl") include("particles.jl") include("particle_advection.jl") +include("forcing_drag.jl") # VERTICAL LEVELS include("vertical_coordinates.jl") @@ -54,7 +55,6 @@ include("stochastic_physics.jl") include("run_speedy.jl") # EXTENSION -include("extending.jl") include("callbacks.jl") include("schedule.jl")