Merge pull request #11 from SpeedyWeather/mk/speedyweather0.9

Adapted to upcoming SpeedyWeather v0.9

Author: milankl

Project.toml

@@ -1,15 +1,15 @@
 name = "StochasticStir"
 uuid = "c3541e7a-ece5-4a49-9931-82ce8f5cb0be"
 authors = ["Milan Klöwer <[email protected]> and contributors"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
 
 [compat]
 DocStringExtensions = "0.9"
-SpeedyWeather = "0.8"
+SpeedyWeather = "0.9"
 julia = "1.8"
 
 [extras]

src/jet_drag.jl

@@ -3,7 +3,7 @@ Defines the JetDrag drag term for SpeedyWeather that applies
 a zonal drag to the BarotropicModel or ShallowWaterModel following
 a jet at specified latitude, strength and width with time scale time_scale.
 $(TYPEDFIELDS)"""
-Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag{NF}
+Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag
 
     # DIMENSIONS from SpectralGrid
     "Spectral resolution as max degree of spherical harmonics"
@@ -14,25 +14,25 @@ Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag{NF}
     time_scale::Second = Day(6)
 
     "Jet strength [m/s]"
-    u₀::Float64 = 20
+    u₀::NF = 20
 
     "latitude of Gaussian jet [˚N]"
-    latitude::Float64 = 30
+    latitude::NF = 30
 
     "Width of Gaussian jet [˚]"
-    width::Float64 = 6
+    width::NF = 6
 
     # TO BE INITIALISED
     "Relaxation back to reference vorticity"
-    ζ₀::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1)
+    ζ₀::LTM{Complex{NF}} = zeros(LTM{Complex{NF}}, trunc+2, trunc+1)
 end
 
 """
 $(TYPEDSIGNATURES)
 Generator function of a JetDrag struct taking resolution and number format from 
 a SpectralGrid. Other options are provided as keyword arguments."""
-function JetDrag(SG::SpectralGrid;kwargs...)
-    return JetDrag{SG.NF}(;SG.trunc,kwargs...)
+function JetDrag(SG::SpectralGrid; kwargs...)
+    return JetDrag{SG.NF}(; SG.trunc, kwargs...)
 end
 
 """
@@ -45,21 +45,21 @@ function SpeedyWeather.initialize!( drag::JetDrag,
                                     model::ModelSetup)
 
     (;spectral_grid, geometry) = model
-    (;Grid,NF,nlat_half) = spectral_grid
+    (;Grid, NF, nlat_half) = spectral_grid
     lat = geometry.latds        # latitude in ˚N for every grid point
 
     # zonal velocity of the jet, allocate following grid and number format NF of model
-    u = zeros(Grid{NF},nlat_half)
+    u = zeros(Grid{NF}, nlat_half)
 
     # Gaussian in latitude, calculated for every grid point ij
     for ij in eachindex(u)
         u[ij] = drag.u₀ * exp(-(lat[ij]-drag.latitude)^2/(2*drag.width^2))
     end
 
     # to spectral space of size lmax+1 × mmax as required by the curl 
-    û = SpeedyTransforms.spectral(u,one_more_degree=true)
+    û = SpeedyTransforms.spectral(u, one_more_degree=true)
     v̂ = zero(û)
-    SpeedyTransforms.curl!(drag.ζ₀,û,v̂,model.spectral_transform)
+    SpeedyTransforms.curl!(drag.ζ₀, û, v̂, model.spectral_transform)
     return nothing
 end
 
@@ -71,7 +71,7 @@ function SpeedyWeather.drag!(
     time::DateTime,
     model::ModelSetup,
 )
-    SpeedyWeather.drag!(diagn,progn,drag,model.geometry)
+    SpeedyWeather.drag!(diagn, progn, drag, model.geometry)
 end
 
 """
@@ -85,16 +85,16 @@ function SpeedyWeather.drag!(
     drag::JetDrag,
     geometry::Geometry,
 )
-    (;vor) = progn
-    (;vor_tend) = diagn.tendencies
-    (;ζ₀) = drag
+    (; vor) = progn
+    (; vor_tend) = diagn.tendencies
+    (; ζ₀) = drag
 
     # 1/time_scale [1/s] but scaled with radius as is the vorticity equation
-    (;radius) = geometry
+    (; radius) = geometry
     r = radius/drag.time_scale.value
 
     # add the -r*(ζ-ζ₀) term to the vorticity tendency that already includes the forcing term
-    @inbounds for lm in eachindex(vor,vor_tend,ζ₀)
+    @inbounds for lm in eachindex(vor, vor_tend, ζ₀)
         vor_tend[lm] -= r*(vor[lm] - ζ₀[lm])
     end
 end

src/stochastic_stirring.jl

@@ -1,7 +1,9 @@
+const LTM = LowerTriangularMatrix   # for convenience
+
 """A SpeedyWeather forcing term that stochastically stirrs relative vorticity in
 the BarotropicModel or the ShallowWaterModel. 
 $(TYPEDFIELDS)"""
-Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF}
+Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing
 
     # DIMENSIONS from SpectralGrid
     "Spectral resolution as max degree of spherical harmonics"
@@ -16,13 +18,13 @@ Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF}
     decorrelation_time::Second = Day(2)
 
     "Stirring strength A [1/s²]"
-    strength::Float64 = 7e-11
+    strength::NF = 7e-11
 
     "Stirring latitude [˚N]"
-    latitude::Float64 = 45
+    latitude::NF = 45
 
     "Stirring width [˚]"
-    width::Float64 = 24
+    width::NF = 24
 
     "Minimum degree of spherical harmonics to force"
     lmin::Int = 8
@@ -38,7 +40,7 @@ Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF}
 
     # TO BE INITIALISED
     "Stochastic stirring term S"
-    S::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1)
+    S::LTM{Complex{NF}} = zeros(LTM{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))
@@ -54,10 +56,10 @@ end
 $(TYPEDSIGNATURES)
 Generator function for StochasticStirring using resolution and number format
 from a SpectralGrid. Further options should be provided as keyword arguments."""
-function StochasticStirring(SG::SpectralGrid;kwargs...)
-    (;trunc,Grid,nlat_half) = SG
-    nlat = RingGrids.get_nlat(Grid,nlat_half)
-    return StochasticStirring{SG.NF}(;trunc,nlat,kwargs...)
+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
 
 """
@@ -70,7 +72,7 @@ function SpeedyWeather.initialize!( forcing::StochasticStirring,
                                     model::ModelSetup)
 
     # precompute forcing strength, scale with radius^2 as is the vorticity equation
-    (;radius) = model.spectral_grid
+    (; radius) = model.spectral_grid
     A = radius^2 * forcing.strength
 
     # precompute noise and auto-regressive factor, packed in RefValue for mutability
@@ -91,12 +93,14 @@ function SpeedyWeather.initialize!( forcing::StochasticStirring,
 end
 
 # function barrier to unpack from model what's needed (spectral transform only here)
-function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
-                                progn::PrognosticVariablesLayer,
-                                forcing::StochasticStirring,
-                                time::DateTime,
-                                model::ModelSetup)
-    SpeedyWeather.forcing!(diagn,forcing,model.spectral_transform)
+function SpeedyWeather.forcing!(
+    diagn::DiagnosticVariablesLayer,
+    progn::PrognosticVariablesLayer,
+    forcing::StochasticStirring,
+    time::DateTime,
+    model::ModelSetup,
+)
+    SpeedyWeather.forcing!(diagn, forcing, model.spectral_transform)
 end
 
 """
@@ -106,9 +110,11 @@ coefficients S of the forcing following an AR1 process in time, transforms
 to grid-point space to apply a mask to only force specified latitudes then
 transforms back to force in spectral space where also the time stepping is
 applied."""
-function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
-                                forcing::StochasticStirring{NF},
-                                spectral_transform::SpectralTransform) where NF
+function SpeedyWeather.forcing!(
+    diagn::DiagnosticVariablesLayer,
+    forcing::StochasticStirring{NF},
+    spectral_transform::SpectralTransform
+) where NF
 
     # noise and auto-regressive factors
     a = forcing.a[]    # = sqrt(1 - exp(-2dt/τ))
@@ -130,14 +136,14 @@ function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
 
     # to grid-point space
     S_grid = diagn.dynamics_variables.a_grid        # reuse general work array
-    SpeedyTransforms.gridded!(S_grid,S,spectral_transform)
+    SpeedyTransforms.gridded!(S_grid, S, spectral_transform)
 
     # mask everything but mid-latitudes
-    RingGrids._scale_lat!(S_grid,forcing.lat_mask)
+    RingGrids._scale_lat!(S_grid, forcing.lat_mask)
 
     # back to spectral space, write directly into vorticity tendency
-    (;vor_tend) = diagn.tendencies
-    SpeedyTransforms.spectral!(vor_tend,S_grid,spectral_transform)
+    (; vor_tend) = diagn.tendencies
+    SpeedyTransforms.spectral!(vor_tend, S_grid, spectral_transform)
 
     return nothing
 end

test/runtests.jl

@@ -3,23 +3,23 @@ using StochasticStir
 using Test
 
 @testset "StochasticStir.jl" begin
-    spectral_grid = SpectralGrid(trunc=31,nlev=1)
+    spectral_grid = SpectralGrid(trunc=31, nlev=1)
 
-    drag = JetDrag(spectral_grid,time_scale=Day(6))
+    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)
+    model = BarotropicModel(;spectral_grid, initial_conditions, forcing, drag)
     simulation = initialize!(model)
 
-    run!(simulation,period=Day(5))
+    run!(simulation, period=Day(5))
     @test simulation.model.feedback.nars_detected == false
 
     # with shallow water model
-    model = ShallowWaterModel(;spectral_grid,initial_conditions,forcing,drag)
+    model = ShallowWaterModel(;spectral_grid, initial_conditions, forcing, drag)
     simulation = initialize!(model)
 
-    run!(simulation,period=Day(5))
+    run!(simulation, period=Day(5))
     @test simulation.model.feedback.nars_detected == false
 end