Rossby-Haurwitz wave initial conditions for eta (#604)

* adds height to Rossby-Haurwitz initial condition

* GeoMakie extension and interactive 3D grid visualisation on the sphere (#600)

* make Geometry parametric with Grid

* GeoMakie extension, globe for interactive 3D visualisation

* add Geodesy as weak dependency too

* ext deps as markdown list

* function args typo

* export globe

* changelog updated

* update docs toml

* remove Geodesy dependency

* move ext into folder

* _faces typo

* add OctahedralGaussian method for _faces

* remove equator for even rings

* include octahedral clenshaw grids

* faces, facesr not returned

* get_vertices for RingGrids

* globe options

* vertices for full grids just rectangles, not diamonds

* globe for data

* extend not overwrite globe

* import Polygon and docstrings

* import Polygon through GeoMakie

* interactive bool

* title corrected

* change vertices definiton to ESWN

* use globe function in docs

* Update Project.toml

Co-authored-by: Anshul Singhvi <[email protected]>

* more globe calls in docs

* docs typo

* Update ext/SpeedyWeatherGeoMakieExt/faces.jl

Co-authored-by: Anshul Singhvi <[email protected]>

---------

Co-authored-by: Anshul Singhvi <[email protected]>

* uses proper model variables for planet independent Rossby-Haurwitz waves

* makes filtering the height anomaly shallow water exclusive

* adds documentation draft for Rossby-Haurwitz wave

* slight changes to the docs

* renames variable for consistency

---------

Co-authored-by: Milan Klöwer <[email protected]>
Co-authored-by: Anshul Singhvi <[email protected]>

CHANGELOG.md

@@ -13,6 +13,7 @@
 - SpectralFilter for horizontal diffusion [#601](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/601)
 - GeoMakie weak dependency, globe function for 3D data visualisation [#600](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/600)
 - Zonal mean for AbstractGridArray [#603](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/603)
+- Rossby-Haurwitz wave with initial conditions for interface displacement for shallow water models[#604](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/604)
 
 ## v0.12.0
 

docs/src/initial_conditions.md

@@ -35,7 +35,7 @@ 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.
+the initial conditions will be set back to what is 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:
 
@@ -47,7 +47,7 @@ barotropic vorticity model) as
 ζ(λ, θ) = 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}``.
+are zonal wavenumber (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
@@ -66,7 +66,7 @@ and (2) To generalise to vertical coordinates, the function `ζ(λ, θ, σ)` tak
 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
+One may 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
@@ -123,16 +123,68 @@ scalars, like `set!(simulation, div=0)` (which is always the case in the
 variables in the spectral space. See `?set!`, the `set!` docstring for more
 details.
 
+## Rossby-Haurwitz wave in a ShallowWater model
+
+For the shallow water model Williamson et al., 1992[^Williamson92] also give 
+initial conditions for the prognostic variable height `h = h_0 + \eta` (equivalent to geopotential).
+The layer thickness is ``h_0`` and ``\eta`` is the interface displacement
+of that layer. SpeedyWeather however, uses as prognostic variable ``\eta``
+for which the initial conditions are
+
+```math
+\begin{align}
+η(λ, θ) &= \frac{R^2}{g} \left( A(θ) + B(θ) \cos(mλ) + C(θ) \cos(2mλ) \right), \\
+
+A(θ) &= \frac{ω(2Ω + ω)}{2}\cos(θ)^2 + \frac{1}{4}K^2\cos(θ)^{2m}\left((m+1)\cos(θ)^2 + (2m^2 - m - 2) - \frac{2m^2}{\cos(θ)^2}\right), \\
+
+B(θ) &= \frac{2(Ω + ω)K}{(m+1)(m+2)} \cos(θ)^m\left((m^2 + 2m + 2) - (m+1)^2\cos(θ)^2\right), \\
+
+C(θ) &= \frac{1}{4}K^2 \cos(θ)^{2m}\left((m+1)\cos(θ)^2 - (m + 2)\right).
+
+\end{align}
+```
+
+Where ``R`` is the radius of the planet on which we consider the
+Rossby-Haurwitz wave, this value can be found in `model.spectral_grid.radius`
+and ``Ω`` and ``g`` are the rotation and the gravity constant of the planet,
+which can be found in `model.planet.rotation` and `model.planet.gravity`.
+
+The interface displacement ``\eta`` in SpeedyWeather's `ShallowWaterModel`
+is stored in the variable `pres` in analogy to the actual pressure in
+the `PrimitiveEquation` model. So we can set ``\eta`` using
+`set!(simulation, pres=η)` for an appropriate implementation of the above
+equations.
+
+With the following we can do a test run of the Rossby-Haurwitz wave in the
+shallow water model without any influences from orography.
+
+```@example haurwitz
+spectral_grid = SpectralGrid(trunc=63, nlayers=1)
+initial_conditions = RossbyHaurwitzWave()
+orography = NoOrography(spectral_grid)
+model = ShallowWaterModel(; spectral_grid, initial_conditions, orography)
+simulation = initialize!(model)
+run!(simulation, period=Day(8))
+
+vor = simulation.diagnostic_variables.grid.vor_grid[:, 1]
+heatmap(vor, title="Relative vorticity [1/s], shallow water Rossby Haurwitz wave after 8 days")
+save("haurwitz_sw.png", ans) # hide
+nothing # hide
+```
+
+There is a noticable difference from the result in the barotropic model, where
+the wave moves around the globe keeping its shape. Here, it deforms around the
+poles and the vorticity patches develop an internal structure.
+
+![Rossby-Haurwitz wave in the shallow water model](haurwitz_sw.png)
+
 ## 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.
+not the initial conditions for other variables.
 
 The following shows that you can set the same `RossbyHaurwitzWave` initial
 conditions also in a `PrimitiveDryModel` (or `Wet`) but you probably
@@ -156,7 +208,7 @@ 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 )
+(`[:, 8]` is the lowermost layer)
 
 ```@example haurwitz
 vor = simulation.diagnostic_variables.grid.vor_grid[:, 8]
@@ -173,4 +225,4 @@ 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)
+[^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)

src/dynamics/initial_conditions.jl

@@ -315,18 +315,34 @@ function initialize!(
 )
     (; m, ω, K, c) = initial_conditions
     (; geometry) = model
+    Ω = model.planet.rotation
+    R = model.spectral_grid.radius
+    g = model.planet.gravity
 
     # 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*λ)
+    # see Williamson et al. 1992, J Computational Physics, eq 147 - 149, 146
+    A(λ, θ) = ω/2 * (2Ω+ω)*cosd(θ)^2 + K^2/4*cosd(θ)^(2m)*((m+1)*cosd(θ)^2 + (2m^2-m-2) - 2m^2/(cosd(θ)^2))
+    B(λ, θ) = (2(Ω+ω)*K)/((m+1)*(m+2))*cosd(θ)^m*( (m^2+2m+2) - (m+1)^2*cosd(θ)^2 )
+    C(λ, θ) = K^2/4*cosd(θ)^(2m)*((m+1) * cosd(θ)^2 - (m + 2))
+
+    η(λ, θ) = R^2/g*(A(λ,θ) + B(λ,θ)*cosd(m*λ) + C(λ,θ)*cosd(2m*λ))
+
     set!(progn, geometry, vor = ζ)
+    model isa ShallowWater && set!(progn, geometry, pres = η)
     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
+    if model isa ShallowWater
+        pres = progn.pres[1]
+        low_value = abs.(pres) .< c
+        pres[low_value] .= 0
+    end
 
     return nothing
 end
@@ -616,4 +632,4 @@ function initialize!(   progn::PrognosticVariables{NF},
     set!(progn, model; pres=η)
 
     return nothing
-end
+end