Merge pull request #5 from SpeedyWeather/mk/docstrings

add docstrings and DocStringExtensions

Author: milankl

Project.toml

@@ -1,14 +1,16 @@
 name = "StochasticStir"
 uuid = "c3541e7a-ece5-4a49-9931-82ce8f5cb0be"
 authors = ["Milan Klöwer <[email protected]> and contributors"]
-version = "0.1"
+version = "0.1.0"
 
 [deps]
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
 
 [compat]
-julia = "1.8"
+DocStringExtensions = "0.9"
 SpeedyWeather = "0.8"
+julia = "1.8"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/StochasticStir.jl

@@ -1,6 +1,7 @@
 module StochasticStir
 
 using SpeedyWeather
+using DocStringExtensions
 
 export StochasticStirring, JetDrag
 

src/jet_drag.jl

@@ -1,3 +1,8 @@
+"""
+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}
 
     # DIMENSIONS from SpectralGrid
@@ -22,42 +27,74 @@ Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag{NF}
     ζ₀::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{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...)
 end
 
+"""
+$(TYPEDSIGNATURES)
+Extends the initialize! function for a JetDrag. Precomputes the zonal velocity
+of the jet, transforms to spectral and takes the curl to get the corresponding
+relative vorticity of the jet, which is used to apply the drag to the vorticity
+equation."""
 function SpeedyWeather.initialize!( drag::JetDrag,
                                     model::ModelSetup)
 
     (;spectral_grid, geometry) = model
     (;Grid,NF,nlat_half) = spectral_grid
-    u = zeros(Grid{NF},nlat_half)
+    lat = geometry.latds        # latitude in ˚N for every grid point
 
-    lat = geometry.latds
+    # zonal velocity of the jet, allocate following grid and number format NF of model
+    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)
     v̂ = zero(û)
     SpeedyTransforms.curl!(drag.ζ₀,û,v̂,model.spectral_transform)
     return nothing
 end
 
-function SpeedyWeather.drag!(   diagn::DiagnosticVariablesLayer,
-                                progn::PrognosticVariablesLayer,
-                                drag::JetDrag,
-                                time::DateTime,
-                                model::ModelSetup)
+# function barrier to unpack only what's needed
+function SpeedyWeather.drag!( 
+    diagn::DiagnosticVariablesLayer,
+    progn::PrognosticVariablesLayer,
+    drag::JetDrag,
+    time::DateTime,
+    model::ModelSetup,
+)
+    SpeedyWeather.drag!(diagn,progn,drag,model.geometry)
+end
 
+"""
+$(TYPEDSIGNATURES)
+Extends the drag! function for JetDrag. Adds the -r*(ζ-ζ₀) term
+to the vorticity tendency.
+"""
+function SpeedyWeather.drag!(
+    diagn::DiagnosticVariablesLayer,
+    progn::PrognosticVariablesLayer,
+    drag::JetDrag,
+    geometry::Geometry,
+)
     (;vor) = progn
     (;vor_tend) = diagn.tendencies
     (;ζ₀) = drag
 
-    (;radius) = model.spectral_grid
+    # 1/time_scale [1/s] but scaled with radius as is the vorticity equation
+    (;radius) = geometry
     r = radius/drag.time_scale.value
-    for lm in eachindex(vor,vor_tend,ζ₀)
+
+    # add the -r*(ζ-ζ₀) term to the vorticity tendency that already includes the forcing term
+    @inbounds for lm in eachindex(vor,vor_tend,ζ₀)
         vor_tend[lm] -= r*(vor[lm] - ζ₀[lm])
     end
 end

src/stochastic_stirring.jl

@@ -1,3 +1,6 @@
+"""A SpeedyWeather forcing term that stochastically stirrs relative vorticity in
+the BarotropicModel or the ShallowWaterModel. 
+$(TYPEDFIELDS)"""
 Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF}
 
     # DIMENSIONS from SpectralGrid
@@ -47,12 +50,22 @@ Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF}
     lat_mask::Vector{NF} = zeros(NF,nlat)
 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...)
 end
 
+"""
+$(TYPEDSIGNATURES)
+Extends the initialize! function for StochasticStirring to precompute
+the AR1 coefficients for the stochastic processes per spherical harmonic.
+Also precomputes a Gaussian latitudinal mask to only force at given latitude
+with width as specified in StochasticStirring."""
 function SpeedyWeather.initialize!( forcing::StochasticStirring,
                                     model::ModelSetup)
 
@@ -61,14 +74,14 @@ function SpeedyWeather.initialize!( forcing::StochasticStirring,
     A = radius^2 * forcing.strength
 
     # precompute noise and auto-regressive factor, packed in RefValue for mutability
-    dt = model.time_stepping.Δt_sec
+    dt = model.time_stepping.Δt_sec         # in seconds
     τ = forcing.decorrelation_time.value    # in seconds
     forcing.a[] = A*sqrt(1 - exp(-2dt/τ))
     forcing.b[] = exp(-dt/τ)
 
-    # precompute the latitudinal mask
+    # precompute the Gaussian latitudinal mask
     (;Grid,nlat_half) = model.spectral_grid
-    latd = RingGrids.get_latd(Grid,nlat_half)
+    latd = RingGrids.get_latd(Grid,nlat_half)   # in ˚N
 
     for j in eachindex(forcing.lat_mask)
         # Gaussian centred at forcing.latitude of width forcing.width
@@ -78,6 +91,7 @@ function SpeedyWeather.initialize!( forcing::StochasticStirring,
     return nothing
 end
 
+# function barrier to unpack from model what's needed (spectral transform only here)
 function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
                                 progn::PrognosticVariablesLayer,
                                 forcing::StochasticStirring,
@@ -86,6 +100,13 @@ function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
     SpeedyWeather.forcing!(diagn,forcing,model.spectral_transform)
 end
 
+"""
+$(TYPEDSIGNATURES)
+Extends the forcing! function for StochasticStirring. Evolves the stochastic
+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
@@ -108,16 +129,15 @@ function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
     end
 
     # to grid-point space
-    S_grid = diagn.dynamics_variables.a_grid
+    S_grid = diagn.dynamics_variables.a_grid        # reuse general work array
     SpeedyTransforms.gridded!(S_grid,S,spectral_transform)
 
     # mask everything but mid-latitudes
     RingGrids._scale_lat!(S_grid,forcing.lat_mask)
 
-    # back to spectral space
+    # back to spectral space, write directly into vorticity tendency
     (;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