SpectralAR1Process and AbstractRandomProcess (#592)

* SpectralAR1Process and AbstractRandomProcess

* update changelog

* netcdf output for random patterns

* random process following ECMWF's SPPT

* Test std and seed of random process, pattern with zero mean

* include in runtests.jl

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Random processes for random pattern generation [#592](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/592)
 - Also allow SpectralGrid as positional argument to model constructors [#593](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/593)
 - De-interweave SpectralTransform [#587](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/587)
 - Rossby-Haurwitz wave initial conditions [#591](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/591)

README.md

@@ -7,18 +7,20 @@
 [![docs](https://img.shields.io/badge/documentation-main-blue.svg)](https://speedyweather.github.io/SpeedyWeather.jl/dev/)
 
 SpeedyWeather.jl is a global spectral atmospheric model with simple physics which is developed as a research playground
-with an everything-flexible attitude as long as it is speedy. With minimal code redundancies it supports
+with an everything-flexible attitude as long as it is speedy. It is easy to use and easy to extend, making 
+atmospheric modelling an interactive experience -- in the terminal, in a notebook or conventionally through scripts.
+With minimal code redundancies it supports
 
 **Dynamics and physics**
-- Different physical equations (barotropic vorticity, shallow water, primitive equations)
+- Different physical equations (barotropic vorticity, shallow water, primitive equations, with and without humidity)
 - Particle advection in 2D for all equations
 - Physics parameterizations for convection, precipitation, boundary layer, etc.
 
 **Numerics and computing**
 - Different spatial grids (full and octahedral grids, Gaussian and Clenshaw-Curtis, HEALPix, OctaHEALPix)
 - Different resolutions (T31 to T1023 and higher, i.e. 400km to 10km using linear, quadratic or cubic truncation)
 - Different arithmetics: Float32 (default), Float64, and (experimental) BFloat16, stochastic rounding
-- multi-threading, layer-wise for dynamics, grid point-wise for physics
+- a very fast and flexible spherical harmonics transform library SpeedyTransforms
 
 **User interface**
 - Extensibility: New model components (incl. parameterizations) can be externally defined
@@ -30,9 +32,9 @@ and Julia will compile to these choices just-in-time.
 
 For an overview of the functionality and explanation see the
 [documentation](https://speedyweather.github.io/SpeedyWeather.jl/dev).
-But as the documentation always lags behind our full functionality you are encouraged
-to [raise an issue](https://github.com/SpeedyWeather/SpeedyWeather.jl/issues) describing
-what you'd like to use SpeedyWeather for.
+You are always encouraged to [raise an issue](https://github.com/SpeedyWeather/SpeedyWeather.jl/issues)
+(even it is not actually an issue but an idea, a suggestion or really anything)
+describing what you'd like to use SpeedyWeather for. We're keen to help!
 
 ## Vision and roadmap
 
@@ -78,7 +80,7 @@ The interface to SpeedyWeather.jl consist of 5 steps: define the grid, create mo
 construct the model, initialize, run
 
 ```julia
-spectral_grid = SpectralGrid(trunc=31, nlayers=8)          # define resolution
+spectral_grid = SpectralGrid(trunc=31, nlayers=8)       # define resolution
 orography = EarthOrography(spectral_grid)               # create non-default components
 model = PrimitiveWetModel(; spectral_grid, orography)   # construct model
 simulation = initialize!(model)                         # initialize all model components

src/SpeedyWeather.jl

@@ -115,6 +115,7 @@ include("dynamics/scaling.jl")
 include("dynamics/tendencies.jl")
 include("dynamics/hole_filling.jl")
 include("dynamics/particle_advection.jl")
+include("dynamics/random_process.jl")
 
 # PARAMETERIZATIONS
 include("physics/albedo.jl")

src/dynamics/diagnostic_variables.jl

@@ -119,7 +119,9 @@ $TYPEDFIELDS."""
     v_grid          ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers)
     "Logarithm of surface pressure [Pa]"
     pres_grid       ::GridVariable2D = zeros(GridVariable2D, nlat_half)
-
+    "Random pattern controlled by random process [1]"
+    random_pattern  ::GridVariable2D = zeros(GridVariable3D, nlat_half)
+
     # PREVIOUS TIME STEP
     "Absolute temperature [K] at previous time step"
     temp_grid_prev  ::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers)

src/dynamics/prognostic_variables.jl

@@ -96,6 +96,9 @@ export PrognosticVariables
     pres::NTuple{NSTEPS, SpectralVariable2D} =
         ntuple(i -> zeros(SpectralVariable2D, trunc+2, trunc+1), NSTEPS)
 
+    "Random pattern following a random process [1]"
+    random_pattern::SpectralVariable2D = zeros(SpectralVariable2D, trunc+2, trunc+1)
+
     "Ocean variables, sea surface temperature and sea ice concentration"
     ocean::PrognosticVariablesOcean{NF, ArrayType, GridVariable2D} =
         PrognosticVariablesOcean{NF, ArrayType, GridVariable2D}(; nlat_half)
@@ -148,6 +151,7 @@ function Base.show(
     println(io, "├ temp:  T$trunc, $nlayers-layer, $NSTEPS-steps LowerTriangularArray{$NF}")
     println(io, "├ humid: T$trunc, $nlayers-layer, $NSTEPS-steps LowerTriangularArray{$NF}")
     println(io, "├ pres:  T$trunc, 1-layer, $NSTEPS-steps LowerTriangularArray{$NF}")
+    println(io, "├ random_pattern: T$trunc, 1-layer LowerTriangularArray{$NF}")
     println(io, "├┐ocean: PrognosticVariablesOcean{$NF}")
     println(io, "│├ sea_surface_temperature:  $nlat-ring $Grid")
     println(io, "│└ sea_ice_concentration:    $nlat-ring $Grid")

src/dynamics/random_process.jl

@@ -0,0 +1,142 @@
+abstract type AbstractRandomProcess <: AbstractModelComponent end
+
+"""$(TYPEDSIGNATURES)
+General transform for `random processes <: AbstractRandomProcess`.
+Takes the spectral `random_pattern` in the prognostic variables
+and transforms it to spectral space in `diagn.grid.random_pattern`."""
+function SpeedyTransforms.transform!(
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    lf::Integer,
+    random_process::AbstractRandomProcess,
+    spectral_transform::SpectralTransform,
+)
+    grid = diagn.grid.random_pattern
+    spec = progn.random_pattern
+    transform!(grid, spec, spectral_transform)
+
+    if :clamp in fieldnames(typeof(random_process))
+        clamp!(grid, random_process.clamp...)
+    end
+end
+
+export NoRandomProcess
+"""Dummy type for no random process."""
+struct NoRandomProcess <: AbstractRandomProcess end
+NoRandomProcess(::SpectralGrid) = NoRandomProcess()
+
+"""$(TYPEDSIGNATURES)
+`NoRandomProcess` does not need to transform any random pattern from
+spectral to grid space."""
+function SpeedyTransforms.transform!(
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    lf::Integer,
+    random_process::NoRandomProcess,
+    spectral_transform::SpectralTransform,
+)
+    return nothing
+end
+
+initialize!(process::NoRandomProcess, model::AbstractModel) = nothing
+random_process!(progn::PrognosticVariables, process::NoRandomProcess) = nothing
+
+export SpectralAR1Process
+
+"""First-order auto-regressive random process (AR1) in spectral space,
+evolving `wavenumbers` with respectice `time_scales` and `standard_deviations`
+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
+    trunc::Int
+
+    "[OPTION] Time scale of the AR1 process"
+    time_scale::Second = Hour(6)
+
+    "[OPTION] Wavenumber of the AR1 process"
+    wavenumber::Int = 12
+
+    "[OPTION] Standard deviation of the AR1 process"
+    standard_deviation::NF = 1/3
+
+    "[OPTION] Range to clamp values into after every transform into grid space"
+    clamp::NTuple{2, NF} = (-1, 1)
+
+    "[OPTION] Random number generator seed"
+    seed::Int = 123
+
+    "Independent random number generator for this random process"
+    random_number_generator::Random.Xoshiro = Random.Xoshiro(seed)
+
+    "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)
+end
+
+# generator function
+SpectralAR1Process(SG::SpectralGrid; kwargs...) = SpectralAR1Process{SG.NF}(trunc=SG.trunc; kwargs...)
+
+function initialize!(
+    process::SpectralAR1Process,
+    model::AbstractModel,
+)
+    # auto-regressive factor in the AR1 process
+    dt = model.time_stepping.Δt_sec         # in seconds
+    process.autoregressive_factor[] = exp(-dt/Second(process.time_scale).value)
+
+    # noise factors per total wavenumber in the AR1 process
+    k = process.wavenumber
+    a = process.autoregressive_factor[]
+    σ = process.standard_deviation
+
+    # ECMWF Tech Memorandum 598, Appendix 8, eq. 18
+    # TODO *norm_sphere seems to be needed, maybe ECMWF uses another normalization of the harmonics?
+    F₀_denominator = 2*sum([(2l + 1)*exp(-l*(l+1)/(k*(k+1))) for l in 1:process.trunc])
+    F₀ = sqrt(σ^2 * (1-a^2) / F₀_denominator)*model.spectral_transform.norm_sphere
+
+    for l in eachindex(process.noise_factors)       # total wavenumber, but 1-based
+        eigenvalue = l*(l-1)                        # (negative) eigenvalue l*(l+1) but 1-based l->l-1
+
+        # ECMWF Tech Memorandum 598, Appendix 8, eq. 17
+        process.noise_factors[l] = F₀*exp(-eigenvalue/(2k*(k+1)))
+    end
+
+    # set mean of random pattern to zero
+    process.noise_factors[1] = 0
+
+    # reseed the random number generator
+    Random.seed!(process.random_number_generator, process.seed)
+    return nothing
+end
+
+function random_process!(
+    progn::PrognosticVariables,
+    process::SpectralAR1Process{NF},
+) where NF
+
+    (; random_pattern) = progn
+    lmax, mmax = size(random_pattern, OneBased, as=Matrix)  # max degree l, order m of harmonics (1-based)
+
+    a = process.autoregressive_factor[]
+    RNG = process.random_number_generator
+    s = convert(NF, 2/sqrt(2))              # to scale: std(real(randn(Complex))) = √2/2 to 1
+
+    lm = 0
+    @inbounds for m in 1:mmax
+        for l in m:lmax
+            lm += 1
+
+            # draw from independent N(0,1) in real and imaginary parts
+            r = s*randn(RNG, Complex{NF})   # scale to unit variance in real/imaginary
+
+            # ECMWF Tech Memorandum 598, Appendix 8, eq. 14
+            ξ = process.noise_factors[l]
+            random_pattern[lm] *= a         # auto-regressive term
+            random_pattern[lm] += ξ*r       # noise term
+        end
+    end
+end

src/dynamics/tendencies.jl

@@ -746,6 +746,9 @@ function SpeedyTransforms.transform!(
     transform!(u_grid, U, S, unscale_coslat=true)
     transform!(v_grid, V, S, unscale_coslat=true)
 
+    # transform random pattern for random process unless NoRandomProcess
+    transform!(diagn, progn, lf, model.random_process, S)
+
     return nothing
 end
 
@@ -784,6 +787,9 @@ function SpeedyTransforms.transform!(
     transform!(u_grid, U, S, unscale_coslat=true)
     transform!(v_grid, V, S, unscale_coslat=true)
 
+    # transform random pattern for random process unless NoRandomProcess
+    transform!(diagn, progn, lf, model.random_process, S)
+
     return nothing
 end
 
@@ -860,6 +866,11 @@ function SpeedyTransforms.transform!(
         @. u_grid_prev = u_grid
         @. v_grid_prev = v_grid
     end
+
+    # transform random pattern for random process unless NoRandomProcess
+    transform!(diagn, progn, lf, model.random_process, S)
+
+    return nothing
 end
 
 """ 

src/dynamics/time_integration.jl

@@ -175,6 +175,10 @@ function leapfrog!(
         spectral_truncation!(var_tend)
         leapfrog!(var_old, var_new, var_tend, dt, lf, model.time_stepping)
     end
+
+    # evolve the random pattern in time
+    random_process!(progn, model.random_process)
+    return nothing
 end
 
 """

src/models/barotropic.jl

@@ -20,6 +20,7 @@ $(TYPEDFIELDS)"""
     DR,     # <:AbstractDrag,
     PA,     # <:AbstractParticleAdvection,
     IC,     # <:AbstractInitialConditions,
+    RP,     # <:AbstractRandomProcess,
     TS,     # <:AbstractTimeStepper,
     ST,     # <:SpectralTransform{NF},
     IM,     # <:AbstractImplicit,
@@ -40,6 +41,7 @@ $(TYPEDFIELDS)"""
     drag::DR = NoDrag()
     particle_advection::PA = NoParticleAdvection()
     initial_conditions::IC = InitialConditions(Barotropic)
+    random_process::RP = NoRandomProcess()
 
     # NUMERICS
     time_stepping::TS = Leapfrog(spectral_grid)
@@ -73,6 +75,7 @@ function initialize!(model::Barotropic; time::DateTime = DEFAULT_DATE)
     initialize!(model.forcing, model)
     initialize!(model.drag, model)
     initialize!(model.horizontal_diffusion, model)
+    initialize!(model.random_process, model)
 
     # initial conditions
     prognostic_variables = PrognosticVariables(spectral_grid, model)

src/models/primitive_dry.jl

@@ -21,6 +21,7 @@ $(TYPEDFIELDS)"""
     AC,     # <:AbstractAdiabaticConversion,
     PA,     # <:AbstractParticleAdvection,
     IC,     # <:AbstractInitialConditions,
+    RP,     # <:AbstractRandomProcess,
     LS,     # <:AbstractLandSeaMask,
     OC,     # <:AbstractOcean,
     LA,     # <:AbstractLand,
@@ -57,6 +58,7 @@ $(TYPEDFIELDS)"""
     adiabatic_conversion::AC = AdiabaticConversion(spectral_grid)
     particle_advection::PA = NoParticleAdvection()
     initial_conditions::IC = InitialConditions(PrimitiveDry)
+    random_process::RP = NoRandomProcess()
 
     # BOUNDARY CONDITIONS
     orography::OR = EarthOrography(spectral_grid)
@@ -110,6 +112,7 @@ function initialize!(model::PrimitiveDry; time::DateTime = DEFAULT_DATE)
     initialize!(model.coriolis, model)
     initialize!(model.geopotential, model)
     initialize!(model.adiabatic_conversion, model)
+    initialize!(model.random_process, model)
 
     # boundary conditionss
     initialize!(model.orography, model)

src/models/primitive_wet.jl

@@ -21,6 +21,7 @@ $(TYPEDFIELDS)"""
     AC,     # <:AbstractAdiabaticConversion,
     PA,     # <:AbstractParticleAdvection,
     IC,     # <:AbstractInitialConditions,
+    RP,     # <:AbstractRandomProcess,
     LS,     # <:AbstractLandSeaMask,
     OC,     # <:AbstractOcean,
     LA,     # <:AbstractLand,
@@ -63,6 +64,7 @@ $(TYPEDFIELDS)"""
     adiabatic_conversion::AC = AdiabaticConversion(spectral_grid)
     particle_advection::PA = NoParticleAdvection()
     initial_conditions::IC = InitialConditions(PrimitiveWet)
+    random_process::RP = NoRandomProcess()
 
     # BOUNDARY CONDITIONS
     orography::OR = EarthOrography(spectral_grid)
@@ -122,6 +124,7 @@ function initialize!(model::PrimitiveWet; time::DateTime = DEFAULT_DATE)
     initialize!(model.coriolis, model)
     initialize!(model.geopotential, model)
     initialize!(model.adiabatic_conversion, model)
+    initialize!(model.random_process, model)
 
     # boundary conditions
     initialize!(model.orography, model)

src/models/shallow_water.jl

@@ -21,6 +21,7 @@ $(TYPEDFIELDS)"""
     DR,     # <:AbstractDrag,
     PA,     # <:AbstractParticleAdvection,
     IC,     # <:AbstractInitialConditions,
+    RP,     # <:AbstractRandomProcess,
     TS,     # <:AbstractTimeStepper,
     ST,     # <:SpectralTransform{NF},
     IM,     # <:AbstractImplicit,
@@ -42,6 +43,7 @@ $(TYPEDFIELDS)"""
     drag::DR = NoDrag()
     particle_advection::PA = NoParticleAdvection()
     initial_conditions::IC = InitialConditions(ShallowWater)
+    random_process::RP = NoRandomProcess()
 
     # NUMERICS
     time_stepping::TS = Leapfrog(spectral_grid)
@@ -77,6 +79,7 @@ function initialize!(model::ShallowWater; time::DateTime = DEFAULT_DATE)
     initialize!(model.drag, model)
     initialize!(model.horizontal_diffusion, model)
     # model.implicit is initialized in first_timesteps!
+    initialize!(model.random_process, model)
 
     # initial conditions
     prognostic_variables = PrognosticVariables(spectral_grid, model)