Enhance docstrings + update MultilayerQG module description in Docs

docs/src/modules/multilayerqg.md

@@ -13,7 +13,7 @@ The QGPV in each layer is
 \mathrm{QGPV}_j = q_j  + \underbrace{f_0 + \beta y}_{\textrm{planetary PV}} + \delta_{j,n} \underbrace{\frac{f_0 h}{H_n}}_{\textrm{topographic PV}}, \quad j = 1, \dots, n \ .
 ```
 
-where ``q_j`` incorporates the relative vorticity in each layer ``\nabla^2\psi_j`` and the 
+where ``q_j`` incorporates the relative vorticity in each layer ``\nabla^2 \psi_j`` and the 
 vortex stretching terms:
 
 ```math
@@ -33,9 +33,9 @@ In view of the relationships above, when we convert to Fourier space ``q``'s and
 related via the matrix equation
 
 ```math
-\begin{pmatrix} \widehat{q}_{\boldsymbol{k}, 1}\\\vdots\\\widehat{q}_{\boldsymbol{k}, n} \end{pmatrix} =
-\underbrace{\left(-|\boldsymbol{k}|^2 \mathbb{1} + \mathbb{F} \right)}_{\equiv \mathbb{S}_{\boldsymbol{k}}}
-\begin{pmatrix} \widehat{\psi}_{\boldsymbol{k}, 1}\\\vdots\\\widehat{\psi}_{\boldsymbol{k}, n} \end{pmatrix}
+\begin{pmatrix} \widehat{q}_{𝐤, 1}\\\vdots\\\widehat{q}_{𝐤, n} \end{pmatrix} =
+\underbrace{\left(-|𝐤|^2 \mathbb{1} + \mathbb{F} \right)}_{\equiv \mathbb{S}_{𝐤}}
+\begin{pmatrix} \widehat{\psi}_{𝐤, 1}\\\vdots\\\widehat{\psi}_{𝐤, n} \end{pmatrix}
 ```
 
 where
@@ -63,6 +63,16 @@ with
 \partial_x Q_j \equiv \delta_{j,n} \partial_x \eta \ .
 ```
 
+### Helper functions
+
+```@docs
+GeophysicalFlows.MultilayerQG.set_q!
+GeophysicalFlows.MultilayerQG.set_ψ!
+GeophysicalFlows.MultilayerQG.updatevars!
+```
+
+### Diagnostics
+
 The eddy kinetic energy in each layer and the eddy potential energy that corresponds to each 
 fluid interface is computed via `energies()`:
 
@@ -79,13 +89,13 @@ GeophysicalFlows.MultilayerQG.fluxes
 
 ### Implementation
 
-Matrices ``\mathbb{S}_{\boldsymbol{k}}`` as well as ``\mathbb{S}^{-1}_{\boldsymbol{k}}`` are included 
-in `params` as `params.S` and `params.S⁻¹` respectively. Additionally, the background PV gradients 
-``\partial_x Q`` and ``\partial_y Q`` are also included in the `params` as `params.Qx` and `params.Qy`.
-
-You can get ``\widehat{\psi}_j`` from ``\widehat{q}_j`` with `streamfunctionfrompv!(psih, qh, params, grid)`, 
-while to get ``\widehat{q}_j`` from ``\widehat{\psi}_j`` you need to call `pvfromstreamfunction!(qh, psih, params, grid)`.
+Matrices ``\mathbb{S}_{𝐤}`` as well as ``\mathbb{S}^{-1}_{𝐤}`` are included in `params` as 
+`params.S` and `params.S⁻¹` respectively. Additionally, the background PV gradients 
+``\partial_x Q`` and ``\partial_y Q`` are also included in the `params` as `params.Qx` and 
+`params.Qy`.
 
+One can get the ``\widehat{\psi}_j`` from ``\widehat{q}_j`` via 
+`streamfunctionfrompv!(psih, qh, params, grid)`, while the inverse, i.e. obtain ``\widehat{q}_j`` from ``\widehat{\psi}_j``, is done via  `pvfromstreamfunction!(qh, psih, params, grid)`.
 
 The equations of motion are time-stepped forward in Fourier space:
 
@@ -99,17 +109,30 @@ In doing so the Jacobian is computed in the conservative form: ``\mathsf{J}(f,g)
 
 Equations are formulated using $q_j$ as the state variables, i.e., `sol = qh`.
 
-Thus:
+The linear operator is constructed in `Equation`
 
-```math
-\begin{aligned}
-\mathcal{L} & = - \nu k^{2n_\nu} \ , \\
-\mathcal{N}(\widehat{q}_j) & = - \widehat{\mathsf{J}(\psi_j, q_j)} - \widehat{U_j \partial_x Q_j} - \widehat{U_j \partial_x q_j}
- + \widehat{(\partial_y \psi_j)(\partial_x Q_j)} - \widehat{(\partial_x \psi_j)(\partial_y Q_j)} + \delta_{j,n} \mu k^{2} \widehat{\psi}_n\ .
-\end{aligned}
+```@docs
+GeophysicalFlows.MultilayerQG.Equation
+GeophysicalFlows.MultilayerQG.hyperviscosity
 ```
-
 
+The nonlinear terms is computed via
+
+```@docs
+GeophysicalFlows.MultilayerQG.calcN!
+```
+which, in turn, calls 
+
+```@docs
+GeophysicalFlows.MultilayerQG.calcN_advection!
+```
+and
+
+```@docs
+GeophysicalFlows.MultilayerQG.addforcing!
+```
+
+
  ## Examples
 
  - `examples/multilayerqg_2layer.jl`: A script that simulates the growth and equilibration of baroclinic eddy turbulence in the Phillips 2-layer model.

src/multilayerqg.jl

@@ -209,21 +209,41 @@ numberoflayers(::SingleLayerParams) = 1
 # Equations
 # ---------
 
-function hyperdissipation(dev, params, grid)
+"""
+    hyperviscosity(dev, params, grid)
+Returns the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with 
+coefficient ``ν`` for ``n`` fluid layers.
+```math
+L_j = - ν |𝐤|^{2 n_ν}, j = 1, ...,n .
+```
+"""
+function hyperviscosity(dev, params, grid)
   T = eltype(grid)
   L = ArrayType(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params)))
   @. L = - params.ν * grid.Krsq^params.nν
   @views @. L[1, 1, :] = 0
   return L
 end
 
+"""
+    LinearEquation(dev, params, grid)
+Returns the `equation` for a multi-layer quasi-geostrophic problem with `params` and `grid`. 
+The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via 
+`hyperviscosity(dev, params, grid)`. The nonlinear term is computed via function `calcNlinear!()`.
+"""
 function LinearEquation(dev, params, grid)
   L = hyperdissipation(dev, params, grid)
   return FourierFlows.Equation(L, calcNlinear!, grid)
 end
-
+
+"""
+    Equation(dev, params, grid)
+Returns the `equation` for a multi-layer quasi-geostrophic problem with `params` and `grid`. 
+The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via 
+`hyperviscosity(dev, params, grid)`. The nonlinear term is computed via function `calcN!()`.
+"""
 function Equation(dev, params, grid)
-  L = hyperdissipation(dev, params, grid)
+  L = hyperviscosity(dev, params, grid)
   return FourierFlows.Equation(L, calcN!, grid)
 end
 
@@ -319,8 +339,8 @@ end
 """
     calcS!(S, Fp, Fm, nlayers, grid)
 
-Constructs the array S, which consists of nlayer x nlayer static arrays S_kl that relate 
-the q's and ψ's at every wavenumber: q̂_{k, l} = S_kl * ψ̂_{k, l}.
+Constructs the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that 
+relate the ``q̂_j``'s and ``ψ̂_j``'s at every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``.
 """
 function calcS!(S, Fp, Fm, nlayers, grid)
   F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp))
@@ -335,8 +355,8 @@ end
 """
     calcS⁻¹!(S, Fp, Fm, nlayers, grid)
 
-Constructs the array S⁻¹, which consists of nlayer x nlayer static arrays (S_kl)⁻¹ that 
-relate the q's and ψ's at every wavenumber: ψ̂_{k, l} = (S_kl)⁻¹ * q̂_{k, l}.
+Constructs the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` 
+that relate the ``q̂_j``'s and ``ψ̂_j``'s at every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``.
 """
 function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid)
   T = eltype(grid)
@@ -355,6 +375,15 @@ end
 # Solvers
 # -------
 
+"""
+    calcN!(N, sol, t, clock, vars, params, grid)
+Compute the nonlinear term, that is the advection term and the forcing,
+```math
+N(q̂_j) = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j}
+ + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j ,
+```
+by calling `calcN_advection!` and `addforcing!`.
+"""
 function calcN!(N, sol, t, clock, vars, params, grid)
   nlayers = numberoflayers(params)
   calcN_advection!(N, sol, vars, params, grid)
@@ -374,7 +403,11 @@ end
 """
     calcN_advection!(N, sol, vars, params, grid)
 
-Calculates the advection term.
+Compute the advection term and stores it in `N`:
+```math
+N(q̂_j) = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j}
+ + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} .
+```
 """
 function calcN_advection!(N, sol, vars, params, grid)
   @. vars.qh = sol
@@ -456,6 +489,11 @@ function calcN_linearadvection!(N, sol, vars, params, grid)
   return nothing
 end
 
+"""
+    addforcing!(N, sol, t, clock, vars, params, grid)
+When the problem includes forcing, calculate the forcing term ``F̂_j`` and add it to the 
+nonlinear term ``N``.
+"""
 addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing
 
 function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
@@ -561,11 +599,12 @@ is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is
 
 The kinetic energy at the ``j``-th fluid layer is 
 ```math
-\\textrm{KE}_j = \\frac{H_j}{H} \\int \\frac1{2} |\\boldsymbol{\\nabla} \\psi_j|^2 \\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad j = 1, \\dots, n \\ ,
+𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n \\ ,
 ```
-while the potential energy that corresponds to the interface ``j+1/2`` (i.e., interface between the ``j``-th and ``(j+1)``-th fluid layer) is
+while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface 
+between the ``j``-th and ``(j+1)``-th fluid layer) is
 ```math
-\\textrm{PE}_{j+1/2} = \\int \\frac1{2} \\frac{f_0^2}{g'_{j+1/2}} (\\psi_j - \\psi_{j+1})^2 \\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad j = 1, \\dots, n-1 \\ .
+𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2}} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2}} \\sum_{𝐤} |ψ_j - ψ_{j+1}|², \\ j = 1, ..., n-1 \\ .
 ```
 """
 function energies(vars, params, grid, sol)
@@ -612,15 +651,14 @@ verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the tota
 
 The lateral eddy fluxes whithin the ``j``-th fluid layer are
 ```math
-\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j \\, v_j \\, \\partial_y u_j 
-\\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad j = 1, \\dots, n \\ ,
+\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j 
+\\frac{𝖽x 𝖽y}{L_x L_y} , \\  j = 1, ..., n \\ ,
 ```
 while the vertical eddy fluxes at the ``j+1/2``-th fluid interface  (i.e., interface between 
 the ``j``-th and ``(j+1)``-th fluid layer) are
 ```math
-\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f_0^2}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, 
-v_{j+1} \\, \\psi_{j} \\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad 
-j = 1 , \\dots , n-1.
+\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, 
+v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1.
 ```
 """
 function fluxes(vars, params, grid, sol)