Rossby-Haurwitz wave initial conditions +docs

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Rossby-Haurwitz wave initial conditions [#591](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/591)
 - Haversine formula and AbstractSphericalDistance [#588](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/588)
 - Array-agnostic SpectralTransform [#583](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/583)
 - Move CUDA dependency into extension [#586](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/586)

docs/make.jl

@@ -14,6 +14,7 @@ makedocs(
                 "How to run SpeedyWeather"=>"how_to_run_speedy.md",
                 "Examples 2D"=>"examples_2D.md",
                 "Examples 3D"=>"examples_3D.md",
+                "Initial conditions" => "initial_conditions.md",
                 "Analysis"=>"analysis.md",
                 "Tree structure"=>"structure.md",
                 "Particle advection"=>"particles.md",

docs/src/initial_conditions.md

@@ -0,0 +1,176 @@
+# Initial conditions
+
+The following showcases some examples of how to set the initial conditions
+for the prognostic variables in SpeedyWeather.jl. In essence there
+are three ways to do this
+
+1. Change the arrays in `simulation.prognostic_variables`
+2. Use the `set!` function
+3. Set the `initial_conditions` component of a model
+
+where 1 is a rather low-level and largely requires you to directly set the
+complex coefficients of the spherical harmonics (advanced!).
+So the `set!` function builds a convenient interface around 1 such that you
+don't have to know about details of grid or spectral space.
+3 then collects method 1 or 2 (or a combination of both) into a single struct
+to "save" some initial conditions for one or several variables.
+This lets you use predefined (inside SpeedyWeather or externally) initial conditions
+as easy as `initial_conditions = RossbyHaurwitzWave()`.
+Let us illustrate this with some examples where we will refer back those
+methods simply as 1, 2, 3.
+
+## Rossby-Haurwitz wave in a BarotropicModel
+
+We define a `BarotropicModel` of some resolution but keep all its components
+as default
+
+```@example haurwitz
+using SpeedyWeather
+spectral_grid = SpectralGrid(trunc=63, nlayers=1)
+model = BarotropicModel(; spectral_grid)
+simulation = initialize!(model)
+```
+
+Now `simulation.prognostic_variables` contains already some
+initial conditions as defined by `model.initial_conditions` (that's method 3).
+Regardless of what those are, we can still mutate them
+before starting a simulation, but if you (re-)initialize the model,
+the initial conditions will be set back to what's defined in `model.initial_conditions.
+We will illustrate the `set!` function now that conveniently lets you (re)set the
+current state of the prognostic variables:
+
+The Rossby-Haurwitz wave[^Williamson92] is defined as an initial condition for
+vorticity ``\zeta`` (which is the sole prognostic variable in the
+barotropic vorticity model) as
+
+```math
+ζ(λ, θ) = 2ω*\sin(θ) - K*\sin(θ)*\cos(θ)^m*(m^2 + 3m + 2)*\cos(m*λ)
+```
+with longitude ``\lambda`` and latitude ``\theta``. The parameters
+are order ``m = 4``, frequencies ``\omega = 7.848e-6s^{-1}, K = 7.848e-6s^{-1}``.
+Now setting these initial conditions is as simple as
+
+```@example haurwitz
+m = 4
+ω = 7.848e-6
+K = 7.848e-6
+
+ζ(λ, θ, σ) = 2ω*sind(θ) - K*sind(θ)*cosd(θ)^m*(m^2 + 3m + 2)*cosd(m*λ)
+set!(simulation, vor=ζ)
+```
+
+with only two difference from the mathematical notation. (1) SpeedyWeather's
+coordinates are in degrees, so we replaced ``\sin, \cos`` with `sind` and `cosd`;
+and (2) To generalise to vertical coordinates, the function `ζ(λ, θ, σ)` takes
+*exactly* three arguments, with `σ` denoting the vertical [Sigma coordinates](@ref).
+This is important so that we can use the same definition of initial conditions
+for the 2D barotropic vorticity model also for the 3D primitive equations.
+
+Some authors filter out low values of spectral vorticity with some cut-off amplitude
+``c = 10^{-10}``, just to illustrate how you would do this (example for method 1)
+
+```@example haurwitz
+c = 1e-10       # cut-off amplitude
+
+# 1 = first leapfrog timestep of spectral vorticity
+vor = simulation.prognostic_variables.vor[1]
+low_values = abs.(vor) .< c
+vor[low_values] .= 0
+nothing # hide
+```
+which is just treating `vor` as an array of something and tweaking the values within!
+
+Let us illustrate these initial conditions. `set!` will set the initial conditions
+in spectral space, taking care of the transform from the equation defined
+in grid coordinates. So to show vorticity again in grid space we transform
+back
+
+```@example haurwitz
+# [1] for first leapfrog time step, [:, 1] for all values on first layer
+vor = simulation.prognostic_variables.vor[1][:, 1]
+vor_grid = transform(vor)
+
+using CairoMakie
+heatmap(vor_grid, title="Relative vorticity [1/s] of Rossby-Haurwitz wave")
+
+save("haurwitz.png", ans) # hide
+nothing # hide
+```
+![Rossby-Haurwitz wave](haurwitz.png)
+
+That's the Rossby-Haurwitz wave! This wave is supposed to travel
+(without changing its shape) eastward around the globe, so let us run
+a simulation for some days
+
+```@example haurwitz
+run!(simulation, period=Day(3))
+
+# a running simulation always transforms spectral variables
+# so we don't have to do the transform manually but just pull 
+# layer 1 (there's only 1) from the diagnostic variables
+vor = simulation.diagnostic_variables.grid.vor_grid[:, 1]
+
+heatmap(vor, title="Relative vorticity [1/s], Rossby Haurwitz wave after 3 days")
+save("haurwitz_day10.png", ans) # hide
+nothing # hide
+```
+![Rossby-Haurwitz wave after day 10](haurwitz_day10.png)
+
+Looks like before, but shifted eastward! That's the Rossby-Haurwitz wave.
+The `set!` function accepts not just a function as outlined above, but also
+scalars, like `set!(simulation, div=0)` (which is always the case in the
+`BarotropicModel`) or grids, or `LowerTriangularArray`s representing
+variables in the spectral space. See `?set!`, the `set!` docstring for more
+details.
+
+## Rossby-Haurwitz wave in primitive equations
+
+We could apply the same to set the Rossby-Haurwitz for a primitive equation
+model, but we have also already defined `RossbyHaurwitzWave` as
+`<: AbstractInitialConditions` so you can use that directly, regardless
+the model. Note that this definition currently only includes vorticity
+not the initial conditions for other variables. Williamson et al. 1992
+define also initial conditions for height/geopotential to be used
+in the shallow water model (eq. 146-149) -- those are currently not included,
+so the wave may not be as stable as its supposed to be.
+
+The following shows that you can set the same `RossbyHaurwitzWave` initial
+conditions also in a `PrimitiveDryModel` (or `Wet`) but you probably
+also want to set initial conditions for temperature and pressure
+to not start at zero Kelvin and zero pressure. Also no orography,
+and let's switch off all physics parameterizations with `physics=false`.
+
+```@example haurwitz
+spectral_grid = SpectralGrid(trunc=42, nlayers=8)
+initial_conditions = InitialConditions(
+                        vordiv=RossbyHaurwitzWave(),
+                        temp=JablonowskiTemperature(),
+                        pres=PressureOnOrography())
+
+orography = NoOrography(spectral_grid)
+model = PrimitiveDryModel(; spectral_grid, initial_conditions, orography, physics=false)
+simulation = initialize!(model)
+run!(simulation, period=Day(5))
+nothing # hide
+```
+
+Note that we chose a lower resolution here (T42) as we are simulating
+8 vertical layers now too. Let us visualise the surface vorticity
+(`[:, 8]` is on layer )
+
+```@example haurwitz
+vor = simulation.diagnostic_variables.grid.vor_grid[:, 8]
+heatmap(vor, title="Relative vorticity [1/s], primitive Rossby-Haurwitz wave")
+
+save("haurwitz_primitive.png", ans) # hide
+nothing # hide
+```
+![Rossby-Haurwitz wave in primitive equations](haurwitz_primitive.png)
+
+As you can see the actual Rossby-Haurwitz wave is not as stable anymore
+(because those initial conditions are not a stable solution of the primitive equations)
+and so the 3-day integration looks already different from the barotropic model!
+
+## References
+
+[^Williamson92]: DL Williamson, JB Drake, JJ Hack, R Jakob, PN Swarztrauber, 1992. *A standard test set for numerical approximations to the shallow water equations in spherical geometry*, **Journal of Computational Physics**, 102, 1, DOI: [10.1016/S0021-9991(05)80016-6](https://doi.org/10.1016/S0021-9991(05)80016-6)

docs/src/orography.md

@@ -120,7 +120,7 @@ is more convenient, for example you can do
 set!(model, orography=(λ,φ) -> 2000*cosd(φ) + 300*sind(λ) + 100*randn())
 
 using CairoMakie
-heatmap(model.orography.orography, title="Zonal 2000m ridge with noise")
+heatmap(model.orography.orography, title="Zonal 2000m ridge [m] with noise")
 save("orography_set.png", ans) # hide
 nothing # hide
 ```
@@ -148,7 +148,7 @@ initialize!(model.orography, model)                 # initially that reset orogr
 # blow up Hawaii by adding a 4000m peak on a 10˚x10˚ large island
 H, λ₀, φ₀, σ = 4000, 200, 20, 5         # height, lon, lat position, and width
 set!(model, orography=(λ,φ) -> H*exp((-(λ-λ₀)^2 - (φ-φ₀)^2)/2σ^2), add=true)
-heatmap(model.orography.orography, title="Super Hawaii orography")
+heatmap(model.orography.orography, title="Super Hawaii orography [m]")
 save("orography_hawaii.png", ans) # hide
 nothing # hide
 ```
@@ -162,6 +162,7 @@ in the surface geopotential `orography.geopot_surf`
 transform!(orography.geopot_surf, orography.orography, model.spectral_transform)
 orography.geopot_surf .*= model.planet.gravity
 spectral_truncation!(orography.geopot_surf)
+nothing # hide
 ```
 
 In the first line, the surface geopotential is still missing the gravity,
@@ -193,25 +194,31 @@ or spherical distance. Compare
 ```@example orography
 λ₀ = 0         # move "Super Hawaii" mountain onto the prime meridian
 set!(model, orography=(λ,φ) -> H*exp((-(λ-λ₀)^2 - (φ-φ₀)^2)/2σ^2))
-heatmap(model.orography.orography, title="Mountain on prime meridian, cartesian coordinates")
+heatmap(model.orography.orography, title="Mountain [m] on prime meridian, cartesian coordinates")
 save("mountain_cartesian.png", ans) # hide
 nothing # hide
 ```
 ![Mountain in cartesian coordinates](mountain_cartesian.png)
 
 which clearly shows that the mountain is only on the Eastern hemisphere
-(the signal on the western hemisphere is due to the interpolation during plotting)
--- probably not what you wanted. Rewrite this as
+-- probably not what you wanted. Note also that because we did not provide
+`add=true` the orography we set through `set!` overwrites the previous
+orography (`add=false` is the default). Rewrite this as
 
 ```@example orography
 λ₀ = 0         # move "Super Hawaii" mountain onto the prime meridian
-set!(model, orography=(λ,φ) -> H*exp(-spherical_distance((λ,φ), (λ₀,φ₀), radius=1)^2/2σ^2))
-heatmap(model.orography.orography, title="Mountain on prime meridian, spherical distance")
+set!(model, orography=(λ,φ) -> H*exp(-spherical_distance((λ,φ), (λ₀,φ₀), radius=360/2π)^2/2σ^2))
+heatmap(model.orography.orography, title="Mountain [m] on prime meridian, spherical distance")
 save("mountain_spherical.png", ans) # hide
 nothing # hide
 ```
 ![Mountain in spherical coordinates](mountain_spherical.png)
 
+And the mountain also shows up on the western hemisphere! Note that we could have defined
+the width ``\sigma`` of the mountain in meters, or we keep using degrees as before
+but then use `radius = 360/2π` to convert radians into degrees. If you set `radius=1`
+then radians are returned and so we could have defined ``\sigma`` in terms of radians too.
+
 ## Defining a new orography type
 
 You can also define a new orography like we defined `ZonalRidge` or `EarthOrography`.

src/RingGrids/geodesics.jl

@@ -18,7 +18,7 @@ struct Haversine <: AbstractSphericalDistance end
 Haversine formula calculating the great-circle or spherical distance (in meters) on the sphere
 between two tuples of  longitude-latitude points in degrees ˚E, ˚N. Use keyword argument `radius`
 to change the radius of the sphere (default 6371e3 meters, Earth's radius), use `radius=1`
-to return the central angle in radians."""
+to return the central angle in radians or `radius=360/2π` to return degrees."""
 function Haversine(             # functor for Haversine struct
     lonlat1::Tuple,             # point 1 in spherical coordinates
     lonlat2::Tuple;             # point 2

src/dynamics/initial_conditions.jl

@@ -293,6 +293,44 @@ function initialize!(   progn::PrognosticVariables{NF},
     return nothing
 end
 
+export RossbyHaurwitzWave
+
+"""Rossby-Haurwitz wave initial conditions as in Williamson et al. 1992, J Computational Physics
+with an additional cut-off amplitude `c` to filter out tiny harmonics in the vorticity field.
+Parameters are $(TYPEDFIELDS)"""
+@kwdef struct RossbyHaurwitzWave <: AbstractInitialConditions
+    m::Int = 4
+    ω::Float64 = 7.848e-6
+    K::Float64 = 7.848e-6
+    c::Float64 = 1e-10
+end
+
+"""$(TYPEDSIGNATURES)
+Rossby-Haurwitz wave initial conditions as in Williamson et al. 1992, J Computational Physics
+with an additional cut-off amplitude `c` to filter out tiny harmonics in the vorticity field."""
+function initialize!(
+    progn::PrognosticVariables,
+    initial_conditions::RossbyHaurwitzWave,
+    model::AbstractModel,
+)
+    (; m, ω, K, c) = initial_conditions
+    (; geometry) = model
+
+    # Rossby-Haurwitz wave defined through vorticity ζ as a function of
+    # longitude λ, latitude θ (in degrees), sigma level σ (vertically constant though)
+    # see Williamson et al. 1992, J Computational Physics, eq 145
+    ζ(λ, θ, σ) = 2ω*sind(θ) - K*sind(θ)*cosd(θ)^m*(m^2 + 3m + 2)*cosd(m*λ)
+    set!(progn, geometry, vor = ζ)
+    set!(progn, geometry, div = 0)  # technically not needed, but set to zero for completeness
+
+    # filter low values below cutoff amplitude c
+    vor = progn.vor[1]  # 1 = first leapfrog timestep
+    low_values = abs.(vor) .< c
+    vor[low_values] .= 0
+
+    return nothing
+end
+
 export JablonowskiTemperature
 
 """

src/output/netcdf_output.jl

@@ -38,7 +38,7 @@ $(TYPEDFIELDS)"""
 
     "[OPTION] also write restart file if output==true?"
     write_restart::Bool = true
-    pkg_version::VersionNumber = pkgversion(SpeedyWeather)
+    pkg_version::VersionNumber = isnothing(pkgversion(SpeedyWeather)) ? v"0.0.0" : pkgversion(SpeedyWeather)
 
     # WHAT/WHEN OPTIONS
     startdate::DateTime = DateTime(2000, 1, 1)