Add forcing and drag extensions to tests

test/extending.jl

@@ -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

test/runtests.jl

@@ -36,5 +36,8 @@ include("column_variables.jl")
 # INITIALIZATION AND INTEGRATION
 include("run_speedy.jl")
 
+# EXTENSION
+include("extending.jl")
+
 # OUTPUT 
 include("netcdf_output.jl")