Merge branch 'main' into mk/heating

CHANGELOG.md

@@ -3,6 +3,8 @@
 ## Unreleased
 
 - ConvectiveHeating implemented [#639](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/639)
+- 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)
 

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)

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)

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

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

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,21 +19,21 @@ 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})
     vordiv = ZonalWind()
     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

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,