MQGeometry-1.0: a multi-layer quasi-geostrophic solver on non-rectangular geometries

This paper presents MQGeometry, a multi-layer quasi-geostrophic (QG) equation solver for non-rectangular geometries. We advect the potential vorticity (PV) with finite volumes to ensure global PV conservation using a staggered discretization of the PV and stream function (SF). Thanks to this staggering, the PV is defined inside the domain, removing the need to define the PV on the domain boundary. We compute PV fluxes with upwind-biased interpolations whose implicit dissipation replaces the usual explicit (hyper-)viscous dissipation. The discretization presented here does not require tuning of any additional parameter, e.g., additional eddy viscosity. We solve the QG elliptic equation with a fast discrete sine transform spectral solver on rectangular geometry. We extend this fast solver to non-rectangular geometries using the capacitance matrix method. Subsequently, we validate our solver on a vortex-shear instability test case in a circular domain, on a vortex–wall interaction test case, and on an idealized wind-driven double-gyre configuration in an octagonal domain at an eddy-permitting resolution. Finally, we release a concise, efficient, and auto-differentiable PyTorch implementation of our method to facilitate future developments on this new discretization, e.g., machine-learning parameterization or data-assimilation techniques.

1 Introduction

Ocean fluid dynamics offers a hierarchy of models with a trade-off between the richness of the physical phenomena and the dimensionality of the system. On the one end of this hierarchy are the Boussinesq non-hydrostatic equations. These equations describe the evolution of the stratification via temperature and salinity; they model explicitly convective phenomena but they require the evolution of six prognostic variables: the three components of the velocity $\mathit{u},\mathit{v},\mathit{w}$; temperature T and salinity s; and the free surface η.

On the other end of this hierarchy are the multi-layer quasi-geostrophic equations. These equations are based on strong hypotheses: static background stratification as well as hydrostatic and geostrophic balances. In a multi-layer quasi-geostrophic (QG) model, the stream function (SF) ψ and potential vorticity (PV) q are stacked in N isopycnal layers with density ρ_i and reference thickness H_i:

$$\begin{array}{ll}{\displaystyle }& {\displaystyle }\mathit{\psi }=[{\mathit{\psi }}_{\mathrm{1}},\mathrm{\dots },{\mathit{\psi }}_N{]}^{\mathrm{T}},\\ {\displaystyle }& {\displaystyle }\mathit{q}=[q_{\mathrm{1}},\mathrm{\dots },q_N{]}^{\mathrm{T}}.\end{array}$$

The governing equations read

$$\begin{array}{}\text{(1a)}& {\displaystyle }& {\displaystyle }{\partial }_t\mathit{q}+\left(\begin{array}{c}\mathit{u}\\ \mathit{v}\end{array}\right)\cdot {\mathrm{\nabla }}_{\mathrm{h}}\mathit{q}=\mathrm{0},\left(\begin{array}{c}\mathit{u}\\ \mathit{v}\end{array}\right)={\mathrm{\nabla }}_{\mathrm{h}}^{⟂}\mathit{\psi },\text{(1b)}& {\displaystyle }& {\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}\mathit{\psi }-f_{\mathrm{0}}^{\mathrm{2}}A\mathit{\psi }=\mathit{q}-\mathit{\beta }y,\text{(1c)}& {\displaystyle }& {\displaystyle }\mathbf{A}=\left[\begin{array}{cccc}\frac{\mathrm{1}}{H_{\mathrm{1}}{g}_{\mathrm{0}}^{\prime }}+\frac{\mathrm{1}}{H_{\mathrm{1}}{g}_{\mathrm{1}}^{\prime }}& \frac{-\mathrm{1}}{H_{\mathrm{1}}{g}_{\mathrm{1}}^{\prime }}& \cdot & \cdot \\ \frac{-\mathrm{1}}{H_{\mathrm{2}}{g}_{\mathrm{1}}^{\prime }}& \frac{\mathrm{1}}{H_{\mathrm{2}}{g}_{\mathrm{1}}^{\prime }}+\frac{\mathrm{1}}{H_{\mathrm{2}}{g}_{\mathrm{2}}^{\prime }}& \frac{-\mathrm{1}}{H_{\mathrm{2}}{g}_{\mathrm{2}}^{\prime }}& \cdot \\ \cdot & \mathrm{\dots }& \mathrm{\dots }& \mathrm{\dots }\\ \cdot & \cdot & \frac{-\mathrm{1}}{H_n{g}_{\mathrm{n}-\mathrm{1}}^{\prime }}& \frac{\mathrm{1}}{H_n{g}_{\mathrm{n}-\mathrm{1}}^{\prime }}\end{array}\right].\end{array}$$

${\mathrm{\nabla }}_{\mathrm{h}}^{⟂}=[-{\partial }_y,{\partial }_x{]}^{\mathrm{T}}$ stands for the horizontal orthogonal gradient, and ${\mathrm{\Delta }}_{\mathrm{h}}={\partial }_{xx}^{\mathrm{2}}+{\partial }_{yy}^{\mathrm{2}}$ denotes the horizontal Laplacian. $f_{\mathrm{0}}+\mathit{\beta }(y-y_{\mathrm{0}})$ is the Coriolis parameter under beta-plane approximation with the meridional axis center; y₀, $g_{\mathrm{0}}^{\prime }=g$ is the gravitational acceleration; and $g_i^{\prime }=g({\mathit{\rho }}_{i+\mathrm{1}}-{\mathit{\rho }}_i)/{\mathit{\rho }}_i$ denotes the reduced gravities. One can include bottom topography as a constant term on the r.h.s. of the QG elliptic Eq. (1b) (Hogg et al., 2014).

QG equations hence involve a single prognostic variable, the potential vorticity (PV) q, which is advected (Eq. 1a). Despite strong hypotheses, the multi-layer QG equations are a robust approximation of ocean meso-scale non-linear dynamics. They offer a computationally efficient playground to study the meso-scale ocean dynamics and to develop physical parameterizations of the unresolved eddy dynamics (e.g., Marshall et al., 2012; Fox-Kemper et al., 2014; Zanna et al., 2017; Ryzhov et al., 2020; Uchida et al., 2022).

Solving the QG equations requires solving an elliptic equation (Eq. 1b) which relates the SF ψ and the potential vorticity q. On a rectangular domain, it can be easily achieved using a fast spectral solver based on discrete Fourier transform (DFT) for periodic boundary conditions or discrete sine transform (DST) for no-flow boundary conditions. Most of the available open-source QG solvers such as GeophysicalFlows (Constantinou et al., 2021), PyQG, or Q-GCM (Hogg et al., 2014) use such spectral solvers and are therefore limited to rectangular geometries.

There are two major issues when implementing QG models on non-rectangular geometries. The first issue is the fact that spectral elliptic solvers do not apply to non-rectangular domains. In ocean models, one typically uses conjugate gradient (CG) iterative solvers such as BiCGSTAB (Van der Vorst, 1992) to solve elliptic equations on non-rectangular domains, for example, when solving the Poisson equations associated with a rigid-lid constraint (Häfner et al., 2021) or for implicit free-surface computations (Kevlahan and Lemarié, 2022) in primitive equation solvers. These CG iterative solvers are significantly slower than spectral solvers for solving Poisson or Helmholtz equations on evenly spaced rectangular grids (Brown, 2020). Using them in a QG solver would reduce significantly the computational efficiency that makes QG appealing compared to shallow-water (SW) models.

The second issue is the definition of the potential vorticity (PV) on the boundaries. If we discretize the PV q on the same locations as the SF ψ, we have to define the PV on the boundaries. This requires using a partial free-slip/no-slip condition (see, e.g., Hogg et al., 2014) to define ghost points in order to compute the Laplacian of the SF.

In this paper, we present MQGeometry, a new multi-layer QG equation solver that addresses these two issues. This solver uses a new discretization of the multi-layer QG equations based on two main choices. The first choice is to discretize the potential vorticity (PV) and the stream function (SF) on two staggered grids. When doing so, the PV is not defined on the boundary of the domain but in its interior. This solves the issue of defining the PV on the domain boundaries. Moreover, this choice makes it possible to use a finite-volume scheme for the PV advection. It guarantees the global conservation of the PV and leverages a fine control on the PV fluxes. The second choice is to solve the elliptic equation using a fast spectral DST solver combined with the capacitance matrix method (Proskurowski and Widlund, 1976) to handle non-rectangular geometries.

In addition to these two choices, we decide to compute PV fluxes with upwind-biased stencil interpolation to remove the additional (hyper-)viscosity used in most discrete QG models. Doing so, we explore a different approach that is complementary to physical parameterizations of the horizontal momentum closure: a careful choice of the numerical schemes used to discretize the continuous equations. This idea has been initially developed by Boris et al. (1992) in the context of large eddy simulation (LES) models and made popular as implicit LES by Grinstein et al. (2007). It has been successfully tested by Von Hardenberg et al. (2000) to study vortex merging in a single-layer QG model and by Roullet et al. (2012) to study the forced-dissipated three-dimensional QG turbulence in a channel configuration. While Von Hardenberg et al. (2000) and Roullet et al. (2012) used only linear upwind-biased reconstruction, we implement here linear and non-linear weighted essentially non-oscillatory (WENO) reconstructions (Jiang and Shu, 1996; Borges et al., 2008) that are tailored to remove spurious numerical oscillations created by linear reconstructions (Liu et al., 1994).

We implement the proposed discretizations in a concise Python PyTorch code that enables seamless GPU acceleration. It benefits from built-in automatic differentiation to facilitate future development in machine learning or data assimilation. We validate our solver on a vortex-shear instability test case in a closed squared domain, an inviscid vortex–wall interaction, and an idealized wind-driven double-gyre configuration in a non-rectangular configuration.

This paper is organized as follows. In Sect. 2, we describe the resolution of the PV advection equation with finite volume. In Sect. 3, we present our elliptic solver based on discrete sine transform and the capacitance matrix method (Proskurowski and Widlund, 1976). In Sect. 4, we detail the solver implementation. In Sect. 5, we describe the experimental settings to validate our discretization. We conclude and evoke further perspectives in Sect. 6.

2 PV advection with finite volumes

2.1 Staggered discretization of PV and SF

The first ingredient of our method is to use finite volumes to solve PV advection (Eq. 1a). This leads naturally to a staggered discretization for the PV and the SF (see right panel of Fig. 1). With this choice, the PV advection (1a) can be integrated over the whole domain as a transported tracer with a finite-volume formulation. Indeed, if ψ is discretized at the cell vertices, the orthogonal gradient of ψ computed with the standard second-order discretization

$${\mathrm{\nabla }}_{\mathrm{h}}^{⟂}{f}_{i,j}=\left(\begin{array}{c}\frac{f_{i,j}-f_{i,j+\mathrm{1}}}{\mathit{\delta }y}\\ \frac{f_{i+\mathrm{1},j}-f_{i,j}}{\mathit{\delta }x}\end{array}\right)$$

lies in the middle of the cell edges. The horizontal velocity u lives in the middle of the vertical edges while the vertical velocity v lives in the middle of the vertical edges. We thus have to discretize the PV q at the cell centers to solve its advection with finite volumes.

Figure 1 Usual (left) and proposed (right) staggered discretization (bottom) of the prognostic variables (q, ψ, u, v) for the QG system (1).

At the boundary which passes along the cell edges, the condition is simply the non-flux of PV across the walls. This staggering clearly separates the boundary condition associated with the transport of PV from the boundary condition associated with optional friction (partial free-slip or no-slip). The global conservation of PV during advection is ensured up to numerical precision with finite volumes.

In the usual discretization (see Appendix A), the PV must be defined on the boundary, which requires definition of the Laplacian operator Δ_h there. This is problematic since it forces blending of the slip boundary conditions into the definition of the Laplacian, rather than having a slip boundary condition clearly separated from the definition of PV. Moreover, depending on this choice, the material conservation of the PV might not be ensured even though the advection scheme, e.g., Arakawa–Lamb (Arakawa and Lamb, 1981), conserves the PV inside the domain.

2.2 Upwinding of PV fluxes

In a finite-volume formulation, we rewrite the PV advection (1a) as the divergence of a flux:

$$\begin{array}{}\text{(2)}& {\displaystyle }{\partial }_t\mathit{q}=-{\mathrm{\nabla }}_{\mathrm{h}}\cdot \left(\left(\begin{array}{c}\mathit{u}\\ \mathit{v}\end{array}\right)\mathit{q}\right).\end{array}$$

The divergence is discretized with the usual second-order finite-difference operator

$${\mathrm{\nabla }}_{\mathrm{h}}\cdot \left(\begin{array}{c}{u}_{i,j}\\ v_{i,j}\end{array}\right)={\displaystyle \frac{u_{i+\mathrm{1},j}-u_{i,j}}{\mathit{\delta }x}}+{\displaystyle \frac{v_{i,j+\mathrm{1}}-v_{i,j}}{\mathit{\delta }y}}.$$

We hence need to interpolate q on the u and v grid points to compute these PV fluxes. In the context of finite-volume methods, we call these interpolations reconstructions.

Figure 2 Illustration of the upwind-biased reconstruction. Away from boundaries we use a five-point stencil for reconstruction, while close to boundaries, we use a three-point stencil when possible (top-left). At distance 1 from the boundary when the velocity moves away from the boundary (bottom-right), we use a second-order reconstruction with a two-point centered stencil rather than upwind-1 reconstruction to have a reconstruction of the order of 2 at least. For linear reconstructions, the weights c_i are fixed. For non-linear reconstruction, the weights c_i depend on the solution.

Figure 3 One-dimensional advection with constant velocity on a periodic domain over one period. We compare three finite-volume methods using linear, WENO-JS (Jiang and Shu, 1996), and WENO-Z (Borges et al., 2008) five-point upwind reconstructions.

For a good trade-off between stability and accuracy, we use a five-point upwind-biased stencil for reconstruction. Near the boundary we use a three-point upwind-biased stencil or a two-point centered stencil as illustrated in Fig. 2. The ordering of the stencil is given by the sign of the velocity. For instance, $q^i={\sum }_{\mathrm{s}=\mathrm{1}}^{\mathrm{5}}{c}_{\mathrm{s}}{q}_{\mathrm{s}}$ when the velocity is positive, and we have the reverse order when the velocity is negative. Using upwind-biased stencils for reconstruction allows us to remove additional ad hoc (hyper-)viscosity, which is necessary when using centered reconstructions (Lemarié et al., 2015). We decide here to rely solely on the upwinding to handle the potential enstrophy dissipation, which is done implicitly.

We offer two possibilities for reconstructions: linear reconstructions or weighted essentially non-oscillatory (WENO) reconstructions (Liu et al., 1994). The weights c_s are fixed for the linear reconstruction. However, flux computed with linear reconstructions tends to produce spurious numerical oscillations when the field q is not smooth (see Fig. 3).

WENO reconstructions benefit from the essentially non-oscillatory property (Harten, 1984): they are designed to prevent spurious numerical oscillations that occur with linear reconstructions (see Fig. 3). With a WENO scheme, the weights c_s are not fixed: they depend on the value of q on the five-point stencil via smoothness indicators, e.g., first- and second-order derivatives computed with finite difference (Jiang and Shu, 1996). The smoother the q on the stencil, the closer the weights c_s are to the linear weights. One can find a detailed explanation of these non-linear weight computations in Borges et al. (2008).

2.3 Time integration

The convergence and stability properties of WENO schemes derived in Jiang and Shu (1996) require the use of a total variation diminishing (TVD) Runge–Kutta scheme of the order of 3 at least (Shu and Osher, 1988). Here we use a TVD Runge–Kutta scheme of the order of 3 (TVD RK3) for time integration with finite time step δt.

Given the time ordinary differential equation

$${\displaystyle }{\partial }_{t}f=Lf,$$

one can write the TVD RK3 scheme with the following three stages:

$$\begin{array}{ll}{\displaystyle }{f}^{\left(\mathrm{1}\right)}& {\displaystyle }=f^{\left(\mathrm{0}\right)}+\mathit{\delta }tL{f}^{\left(\mathrm{0}\right)},\\ {\displaystyle }{f}^{\left(\mathrm{2}\right)}& {\displaystyle }=f^{\left(\mathrm{1}\right)}+{\displaystyle \frac{\mathit{\delta }t}{\mathrm{4}}}\left(L{f}^{\left(\mathrm{1}\right)}-\mathrm{3}L{f}^{\left(\mathrm{0}\right)}\right),\\ {\displaystyle }{f}^{\left(\mathrm{3}\right)}& {\displaystyle }=f^{\left(\mathrm{2}\right)}+{\displaystyle \frac{\mathit{\delta }t}{\mathrm{12}}}\left(\mathrm{8}L{f}^{\left(\mathrm{2}\right)}-L{f}^{\left(\mathrm{1}\right)}-L{f}^{\left(\mathrm{0}\right)}\right).\end{array}$$

This low-memory version requires only the storage of the intermediate time derivatives Lf⁽ⁱ⁾.

3 Spectral elliptic solver on non-rectangular domain

3.1 Staggering of PV and SF

After having solved PV advection equation (Eq. 1a), one has to solve the elliptic Eq. (1b) to compute the SF. The SF ψ satisfies the homogeneous Dirichlet boundary condition; hence we have to compute the r.h.s. terms inside the domain. Due to the staggered discretization, we need to interpolate them between the two grids. A natural way to proceed is to use a four-point linear interpolation to interpolate q on the ψ grid (see Fig. 4).

Figure 4 Illustration of the four-point interpolation to interpolate q on the ψ grid.

3.2 DST solver on rectangular geometry

To solve the elliptic Eq. (1b), we use a vertical transform and a fast spectral solver with discrete sine transform (DST). One can diagonalize the matrix A (defined in Eq. 1c) as follows:

$${\displaystyle }\mathbf{A}={\mathbf{C}}_{\mathrm{m}\mathrm{2}\mathrm{l}}\mathbf{\Lambda }{\mathbf{C}}_{\mathrm{l}\mathrm{2}\mathrm{m}},$$

where the layer-to-mode matrix C_l2m is the inverse of the mode-to-layer matrix C_m2l, and Λ is a diagonal matrix containing A eigenvalues. One can then perform the following layer-to-mode transform:

$${\displaystyle }\stackrel{\mathrm{̃}}{\mathit{\psi }},\stackrel{\mathrm{̃}}{\mathit{q}}={\mathbf{C}}_{\mathrm{l}\mathrm{2}\mathrm{m}}\mathit{\psi },{\mathbf{C}}_{\mathrm{l}\mathrm{2}\mathrm{m}}\mathit{q}.$$

With this transform, the elliptic Eq. (1b) becomes a stack of N two-dimensional Helmholtz equations

$$\begin{array}{}\text{(3)}& {\mathrm{\Delta }}_{\mathrm{h}}\stackrel{\mathrm{̃}}{\mathit{\psi }}-f_{\mathrm{0}}^{\mathrm{2}}\mathbf{\Lambda }\stackrel{\mathrm{̃}}{\mathit{\psi }}=\stackrel{\mathrm{̃}}{\mathit{q}}-\mathit{\beta }y,\end{array}$$

with homogeneous Dirichlet boundary conditions. To solve these N two-dimensional Helmholtz equations, we use fast diagonalization with type-I discrete sine transform (DST-I) since the usual five-point Laplacian operator

$${\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}{f}_{i,j}={\displaystyle \frac{f_{i+\mathrm{1},j}-\mathrm{2}{f}_{i,j}+f_{i-\mathrm{1},j}}{\mathit{\delta }{x}^{\mathrm{2}}}}+{\displaystyle \frac{f_{i,j+\mathrm{1}}-\mathrm{2}{f}_{i,j}+f_{i,j+\mathrm{1}}}{\mathit{\delta }{y}^{\mathrm{2}}}}$$

becomes a diagonal operator in the type-I sine basis (Press and Teukolsky, 2007, Chap. 20.4, p. 1055). After having solved these N Helmholtz equations, we transform back ψ from mode to layers:

$$\mathit{\psi }={\mathbf{C}}_{\mathrm{m}\mathrm{2}\mathrm{l}}\stackrel{\mathrm{̃}}{\mathit{\psi }}.$$

3.3 Capacitance matrix method for non-rectangular geometry

To handle non-rectangular domains, we use the capacitance matrix method (Proskurowski and Widlund, 1976). We have a non-rectangular domain $\mathrm{\Omega }=\partial \mathrm{\Omega }\cup \stackrel{\circ }{\mathrm{\Omega }}$, where ∂Ω is the domain boundary and $\stackrel{\circ }{\mathrm{\Omega }}$ is the interior of the domain. Our non-rectangular domain Ω is embedded in a rectangular domain ${\mathrm{\Omega }}_{\mathrm{R}}=\partial {\mathrm{\Omega }}_{\mathrm{R}}\cup {\stackrel{\circ }{\mathrm{\Omega }}}_{\mathrm{R}}$. We denote by $\mathcal{I}=\partial \mathrm{\Omega }\setminus \partial {\mathrm{\Omega }}_{\mathrm{R}}$ the set of non-rectangular boundary indices (see Fig. 5). We assume that we have K non-rectangular boundary points I_k∈ℐ.

Figure 5 Domain Ω included in a rectangular domain Ω_R and the set $\mathcal{I}=\partial \mathrm{\Omega }\setminus \partial {\mathrm{\Omega }}_{\mathrm{R}}$ of non-rectangular boundary indices.

We now explain how to solve the following non-rectangular Helmholtz equation:

$$\begin{array}{}\text{(4)}& {\displaystyle }{\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}f-\mathit{\lambda }f=r\text{ in }\stackrel{\circ }{\mathrm{\Omega }},\mathit{\lambda }\in {\mathbb{R}}^{+},f=\mathrm{0}\text{ on }\partial \mathrm{\Omega }.\end{array}$$

One can find a more detailed explanation in Blayo and LeProvost (1993), as this is our inspiration for using the capacitance matrix method.

3.3.1 Precomputations

For each non-rectangular boundary point I_k∈ℐ, a Green function g_k is defined as the solution of the following rectangular Helmholtz equation:

$$\begin{array}{ll}{\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}{g}_k-\mathit{\lambda }{g}_k& {\displaystyle }=\left\{\begin{array}{l}\mathrm{1}\text{ on }{I}_k\\ \mathrm{0}\text{ in }{\stackrel{\circ }{\mathrm{\Omega }}}_{\mathrm{R}}\setminus \mathit{\left\{}{I}_k\mathit{\right\}}\end{array}\right.,\\ {\displaystyle }{g}_k& {\displaystyle }=\mathrm{0}\text{ on }\partial {\mathrm{\Omega }}_{\mathrm{R}},\end{array}$$

solved using DST-I fast diagonalization. With these Green functions, we compute the square matrix M with coefficients

$$m_{k,l}=g_l\left(I_k\right).$$

We compute this matrix inverse to get the so-called capacitance matrix $\mathbf{C}={\mathbf{M}}^{-\mathrm{1}}$.

3.3.2 First step

In the first step, we solve the following rectangular Helmholtz equation:

$$\begin{array}{ll}{\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}{f}^{\left(\mathrm{1}\right)}-\mathit{\lambda }{f}^{\left(\mathrm{1}\right)}& {\displaystyle }=\left\{\begin{array}{l}r\text{ in }\stackrel{\circ }{\mathrm{\Omega }}\\ \mathrm{0}\text{ in }{\stackrel{\circ }{\mathrm{\Omega }}}_{\mathrm{R}}\setminus \stackrel{\circ }{\mathrm{\Omega }}\end{array}\right.\\ {\displaystyle }{f}^{\left(\mathrm{1}\right)}& {\displaystyle }=\mathrm{0}\text{ on }\partial {\mathrm{\Omega }}_{\mathrm{R}},\end{array}$$

using DST-I fast diagonalization. We then compute the vector s coefficients

$${\displaystyle }{s}_k=f^{\left(\mathrm{1}\right)}\left(I_k\right),$$

and we deduce the vector

$${\displaystyle }\mathit{\alpha }=-\mathbf{C}\mathit{s}.$$

3.3.3 Second step

In the second step, we solve the following rectangular Helmholtz equation:

$$\begin{array}{ll}{\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}{f}^{\left(\mathrm{2}\right)}-\mathit{\lambda }{f}^{\left(\mathrm{2}\right)}& {\displaystyle }=\left\{\begin{array}{l}r\text{ in }\stackrel{\circ }{\mathrm{\Omega }}\\ {\mathit{\alpha }}_k\text{ on }{I}_k\\ \mathrm{0}\text{ in }{\stackrel{\circ }{\mathrm{\Omega }}}_{\mathrm{R}}\setminus \mathrm{\Omega },\end{array}\right.\\ {\displaystyle }{f}^{\left(\mathrm{2}\right)}& {\displaystyle }=\mathrm{0}\text{ on }\partial {\mathrm{\Omega }}_{\mathrm{R}},\end{array}$$

using DST-I fast diagonalization. This function f⁽²⁾ is such that

$$\begin{array}{ll}{\displaystyle }{\mathrm{\Delta }}_{\mathrm{h}}{f}^{\left(\mathrm{2}\right)}-\mathit{\lambda }{f}^{\left(\mathrm{2}\right)}& {\displaystyle }=r\text{ in }\stackrel{\circ }{\mathrm{\Omega }}\\ {\displaystyle }{f}^{\left(\mathrm{2}\right)}\left(I_k\right)& {\displaystyle }=\mathrm{0}\forall I_k\in \mathcal{I}\\ {\displaystyle }{f}^{\left(\mathrm{2}\right)}& {\displaystyle }=\mathrm{0}\text{ on }\partial \mathrm{\Omega }\setminus \mathcal{I}.\end{array}$$

The restriction of f⁽²⁾ over our non-rectangular domain Ω is therefore the solution of the Helmholtz Eq. (4).

3.3.4 Numerical cost

The capacitance matrix method involves solving a dense linear problem with K unknowns, K being the number of boundary irregular points ℐ. The capacitance matrix is hence a K×K matrix, and its inversions require 𝒪(K³) computations.

In practice, the inversion is the most computationally expensive operation, but it is precomputed a single time. The major limitation of this method is in terms of memory, i.e., to store the K×K matrix especially on GPUs whose memory capacity is usually smaller than CPUs. On the laptop utilized, we were able to use up to K=10 000, allowing us to run simulations akin to the North Atlantic at a resolution of 6 km. At this resolution, the ageostrophic effects become significant, and the QG hypothesis might no longer be relevant.

4 Implementation

4.1 Programming language and library

One can implement the above discretization using any programming language. To anticipate later applications in data assimilation and machine learning, and to benefit easily from GPU acceleration, we have decided to use the PyTorch library (Paszke et al., 2019). This library contains the usual linear algebra routines, and the computations can be easily vectorized. It offers a built-in automatic differentiation to compute gradients or adjoints. This can be beneficial for future machine-learning and data-assimilation developments.

4.2 Upwind flux computation

Upwind flux computations require splitting the velocity into a positive and a negative part. This can be achieved very efficiently using the ReLU function

$${\displaystyle }\text{ReLU}\left(x\right)=\text{max}(x,\mathrm{0}),$$

whose PyTorch implementation is highly optimized since this function is widely use in neural networks. The positive part u⁺ and negative part u⁻ of the velocity are given by

$$\begin{array}{ll}{\displaystyle }{u}^{+}& {\displaystyle }=\text{ReLU}\left(u\right),\\ {\displaystyle }{u}^{-}& {\displaystyle }=u-u^{+}.\end{array}$$

4.3 WENO reconstructions

Weighted essentially non-oscillatory (WENO; Liu et al., 1994) is a large class of reconstruction methods. We implement two of the most widely used WENO reconstructions: WENO-JS (Jiang and Shu, 1996) and WENO-Z (Borges et al., 2008). Implementing these methods requires only a few lines of Python code (shown in Listing B1 in Appendix B).

4.4 Masks

For the non-rectangular domain, one has to provide a binary mask for the PV grid with ones inside the domain and zeros outside. Given this mask, we automatically compute the masks for the other variable, i.e., the SF ψ and the two components of the velocity u and v, and we deduce the non-rectangular boundary set ℐ involved in the capacitance matrix computations. We also compute specific masks that specify the stencil for PV reconstruction on u and v points for PV fluxes computations. This is illustrated in Fig. 6.

Figure 6 Illustration of the different masks. The label “u ic2” means that one uses two-point centered stencils to reconstruct q on these u points for computations of PV fluxes. The label “v iup5” means that one uses five-point upwind-biased stencils to reconstruct q on these v points for computations of vertical PV fluxes.

4.5 Elliptic solver

Our elliptic solver is based on fast diagonalization using discrete sine transform. Since PyTorch implements FFT but not DST, we implement DST-I using FFTs with specific pre- and post-processing operations. Notably, our DST-I implementation is faster than SciPy's implementation on Intel CPUs thanks to PyTorch bindings to MKL FFT.

The capacitance matrix method implementation is straightforward following the equations in Sect. 3. Capacitance matrices are precomputed given the set of non-rectangular boundary indices ℐ. The method solves the Helmholtz equation with machine precision accuracy. Figure 7 illustrates this point on a circular domain at resolution 256². From a prescribed stream function f, taken as Gaussian white noise and vanishing along the boundary, we apply the Helmholtz operator to get a r.h.s. We then solve the Helmholtz equation with this r.h.s. to get f_inv. Figure 7 shows that f_inv−f is of the order of machine precision. The method is independent of the domain shape. The numerical experiments below explore various domain shapes.

Figure 7 Example of Helmholtz equation solved with our solver on a circular domain embedded in a 256² square domain. The true function f is initialized with spatially uncorrelated Gaussian white noise. The inverted function f_inv is computed using our solver. The difference between f and f_inv is of the order of machine precision. CMM stands for capacitance matrix method.

Our spectral solver cannot handle curvilinear coordinates or non-constant metric terms. Solving the QG equations on a grid with non-constant metric terms requires an iterative elliptic solver, e.g., a multigrid solver (Fulton et al., 1986). Iterative solvers can be sped up with a good initial guess. One can still use the solver presented here to provide an initial guess assuming that the metric terms dx and dy are constant, e.g., equal to the mean dx and dy.

4.6 Compilation

PyTorch embeds a compiler which enables a significant speed-up for computationally intense routines. We compile the flux computations and finite-difference routines. We get a 2.2× speed-up thanks to this compilation with torch.compile.

4.7 Ensemble simulations

Thanks to PyTorch vectorized operations, our implementation allows us to run ensemble simulations, i.e., parallel simulations starting from different initial conditions. This is a promising possibility for later developments of ensemble-based data assimilation or for stochastic extensions such as Li et al. (2023).

4.8 Architecture

The code is divided into six Python scripts:

• helmholtz.py, which contains the Helmholtz solvers based on DST-I and the capacitance matrix method

• fd.py, which contains the finite-difference functions

• reconstruction.py, which contains the reconstruction (i.e., interpolations in the context of finite volumes) routines

• masks.py, which contains the mask utility module

• flux.py, which contains the flux computation routines

• qgm.py, which contains the Python class implementing the multi-layer QG model.

We end up with a concise PyTorch implementation (∼ 750 lines of code) which is as close as possible to the equations. It can run seamlessly on CPUs and GPUs with CUDA compatibility.

4.9 Performance

To assess the performance of our solver, we ran the double-gyre experiment (see below) on a Dell Precision 7560 laptop equipped with an Intel Core i9-11950H CPU and an NVIDIA RTX A3000 laptop GPU. We measure the number of seconds per grid point per time step as well as the power consumption of the devices using the commands nvidia-smi for GPU and turbostat for CPU.

Using the GPU, we get an execution time of $\mathrm{2.5}\times {\mathrm{10}}^{-\mathrm{8}}$ s per grid point per time step with a power consumption of 75 W. Using the CPU with a single core, we get an execution time of $\mathrm{1.9}\times {\mathrm{10}}^{-\mathrm{7}}$ s per grid point per time step with a power consumption of 22 W. Using the full CPU with 16 cores, we get an execution time of $\mathrm{7.0}\times {\mathrm{10}}^{-\mathrm{8}}$ s per grid point per time step with a power consumption of 67 W.

As expected, the GPU has a better power efficiency than the CPU (Häfner et al., 2021). There is still room for improvement for CPU parallelization in our code, as the native parallelism of PyTorch is not fully efficient in our case.

Figure 8 Vortex-shear instability experiment. Evolution of potential vorticity q at initial time t=0; at intermediate times $\mathrm{5}\mathit{\tau },\mathrm{7.5}\mathit{\tau },\mathrm{9}\mathit{\tau },\mathrm{13}\mathit{\tau },\mathrm{18}\mathit{\tau }$, and 24τ; and at final simulation time 30τ.

4.10 Accuracy

Despite the high-order WENO reconstruction used to solve the PV transport, our solver is formally second-order accurate due to the staggering and the use of a second-order perpendicular and divergence operator. This calls for a discussion.

The perpendicular gradient operator is applied to the SF, which is the smoothest field that we resolve. The benefits of using higher-order schemes for this operator are less obvious. The second-order divergence operator used in finite-volume advection ensures the global conservation of PV up to numerical precision. Higher-order schemes might discard this conservation property. For the PV advection, low-order reconstruction schemes suffer from higher numerical diffusion (Lemarié et al., 2015) while linear reconstructions tend to create more oscillations for non-smooth fields. These two considerations motivate the use of high-order non-linear reconstruction on the PV field, which is non-smooth since QG flows contain boundary currents, eddies, and filaments.

5 Numerical validation

We run three numerical experiments involving meso-scale vortices (20–200 km diameters) to validate our solver. The first one is a vortex-shear instability, the second is a vortex–wall interaction, and the third is an idealized double-gyre experiment, which is a usual toy model for western boundary currents. With our solver, multi-layer QG equations deliver on their promise of a computationally efficient playground to study meso-scale non-linear dynamics. Indeed, the three experiments presented here ran on a laptop and took a few to 50 min to run. We provide the Python scripts to reproduce these experiments and the figures.

5.1 Vortex-shear instability

The first validation of our method consists of a vortex-shear instability. We consider a rotating fluid in a circular domain with diameter D=100 km embedded in a square domain of size $L_x\times L_y=$100 km× 100 km on an f plane with a Coriolis parameter f₀, whose value is deduced from the Burger number below. There is a single layer of fluid with reference thickness H=1 km. We assume no-flow and free-slip boundary conditions. The gravity constant is set to g=10 m s⁻².

We study the shear instability of a meso-scale shielded Rankine vortex, which has piece-wise constant PV. In the initial state, the vortex is composed of a core vortex surrounded by a ring of opposite-sign PV to the core such that the total sum of the PV is 0. This system is shear unstable and generates multipoles (Morel and Carton, 1994). We focus here on the tripole formation regime.

The core of the vortex has a radius r₀=10 km and positive vorticity. The surrounding ring has an inner radius r₀=10 km and r₁=14 km. The remaining parameters of the simulation are set via the Rossby and Burger numbers defined as follows:

$$\begin{array}{}\text{(5)}& {\displaystyle }Ro& {\displaystyle }={\displaystyle \frac{u_{\max }}{f_{\mathrm{0}}{r}_{\mathrm{0}}}},\text{(6)}& {\displaystyle }Bu& {\displaystyle }={\displaystyle \frac{gH}{(f_{\mathrm{0}}{r}_{\mathrm{0}}{)}^{\mathrm{2}}}},\end{array}$$

where u_max is the maximum velocity of the initial condition. Given the Burger number, we compute the Coriolis parameter using $f_{\mathrm{0}}=\sqrt{\frac{gH}{Bu{r}_{\mathrm{0}}^{\mathrm{2}}}}$. Then given the Rossby number, we rescale the velocity field of the initial condition such that the maximum velocity is u_max=Rof₀r₀. The quasi-geostrophic equations are valid for Bu≤1 and Ro≪1. We perform the experiment with Ro=0.01 and Bu=1. The equations are integrated over a period of 30 τ, with $\mathit{\tau }=\Vert {\mathit{q}}_{\mathrm{init}}{\Vert }_{\mathrm{2}}^{-\mathrm{1}}$ the eddy-turnover time and q_init the initial condition, shown at the top left of Fig. 8. At initial time the PV contours of the core and the ring r=r_i are slightly perturbed by means of a mode-3 azimuthal perturbation defined in polar coordinates (r,θ) by

$${\displaystyle }\left(\mathrm{1}+\mathit{ϵ}\cos \left(\mathrm{3}\mathit{\theta }\right)\right)r=r_i,$$

where ϵ≪1 is a small parameter, typically ϵ=0.001. This perturbation favors the growth of the most unstable mode which, given the ratio $r_{\mathrm{1}}/r_{\mathrm{0}}$, evolves non-linearly into a tripole (Morel and Carton, 1994).

We use WENO-Z reconstructions for this experiment. This experiment can be reproduced with script vortex_shear.py. We run the reference experiment at a resolution of 1024².

We see in Fig. 8 the evolution of the vortex PV q at different times between the beginning and the end of the simulation at time t=30 τ. At time t=5 τ the result of the instability is visible and the core vortex which was initially circular has a triangular shape. At time t=7.5 τ the outer positive PV ring has become the expected tripole (Morel and Carton, 1994). At time t=9 τ the core vortex recovers a triangular shape with negative PV filaments ranging from the triangle vertices to the three positive vortices. At time t=13 τ, the core vortex keeps a triangular shape and the vorticity filaments become thinner. At time t=18 τ, the filament thickness reaches the grid scale and the filaments start being dissipated by WENO-Z implicit dissipation. At time t=24 τ, the filaments are progressively dissipated and have a smaller amplitude. The core vortex shape becomes more circular, and the three positive vortices are surrounded by a smaller negative vortex. At final time t=30 τ, the filament amplitude has lowered, and small vorticity dipoles appear. The order-3 symmetry, which was injected by the initial perturbation, starts to be lost only at time t=30 τ. For later times (not shown), the system evolves into a chaotic system. This chaotic evolution is expected as the system is sensitive to initial conditions.

Figure 9 Evolution of the total enstrophy for the vortex-shear at resolutions 1024², 512², and 256² using WENO-3 or WENO-5 for PV reconstruction.

To measure the sensitivity of our solver to the resolution and the order of the reconstruction scheme (3 or 5), we compare the solutions produced by our solver at resolutions 1024², 512², and 256² using third-order WENO (WENO-3) and fifth-order WENO (WENO-5) reconstructions. In Fig. 9 we plot the evolution of the total enstrophy at the three different resolutions. The results show that the enstrophy dissipation decreases as the resolution increases, indicating that a higher resolution better preserves the PV variance, and that low-resolution simulations suffer from an excessive dissipation.

We also note that WENO-3 leads to excessive numerical dissipation compared to WENO-5. Moreover, the enstrophy is better preserved with WENO-5 at resolution 512² than with WENO-3 at resolution 1024². We plot in Fig. C1 in the Appendix the final state of the simulation at resolutions 512² and 1024² with WENO-3 and WENO-5. This indicates that WENO-5 increases the effective resolution. In terms of computation cost, the simulation runtime is 1 min 2 s with WENO-5 at resolution 512² and 5 min 52 s with WENO-3 at resolution 1024² with the NVIDIA RTX A3000 laptop GPU. This illustrates the benefits of using high-order reconstructions despite the fact that our code is globally second-order accurate.

5.2 Vortex–wall interaction

To challenge the numerics we now study the propagation of a single meso-scale vortex along a solid boundary, with a free-slip boundary condition. The domain is square with a thin-wall obstacle. Because the flow is inviscid, up to numerical errors, the vortex is expected to follow the boundary according to the mirror effect. In particular the vortex should slip around the obstacle and circumvent it, without detaching from it and without producing filaments (Deremble et al., 2016). The challenge is thus to have a solution that is as inviscid as possible and to enforce as much as possible both the no-flow and the free-slip boundary conditions.

Figure 10 Vortex–wall interaction experiment. (a) Vortex potential vorticity q and stream function ψ contours at initial time t=0; at intermediate times $\mathrm{3.1}\mathit{\tau },\mathrm{6.3}\mathit{\tau },\mathrm{9.4}\mathit{\tau },\mathrm{12.6}\mathit{\tau },\mathrm{15.7}\mathit{\tau }$, and 18.9τ; and at final simulation time 22τ. (b) Evolution of the total enstrophy at resolutions 1024², 512², and 256².

The setup is as follows. The domain size is $L_x\times L_y=$100 km × 100 km on an f plane with a Coriolis parameter f₀. The thin-wall obstacle is vertical, starting from the middle of the domain south boundary, and of length $L_y/\mathrm{4}$ (see Fig. 10) and has a two-cell width. There is a single layer of fluid with reference thickness H=1 km. We assume no-flow and free-slip boundary conditions. The gravity constant is set to g=10 m s⁻². In the initial state, the vortex is a circle of radius r₀=10 km with constant PV and is at the bottom of the domain and at the left of the wall. The circular shape differs from the oval shape a vortex has when moving along a rectilinear wall. The remaining parameters of the simulation are set via the Rossby and Burger numbers with Ro=0.01, and Bu=1.

Figure 10 shows PV snapshots at various times, superimposed with the SF. The vortex behaves according to the inviscid regime. It clearly follows the wall and the obstacle elastically, without any sign of dissipative process such as filament detachment. The vortex at t=22 τ has recovered the characteristic oval shape. Between t=9.4 τ and t=12.6 τ, the vortex circumvents the edge, which causes it to experience its maximal deformation, but the solution remains smooth. The SF is clearly constant along the boundary, as a direct consequence of the capacitance matrix method. During the circumvention, it remains so, despite the thin wall imposing a strong curvature at the edge.

This experiment shows that the numerics has very good conservation properties on inviscid flows. This may come as a surprise since the upwinding does induce a numerical dissipation. In practice, the dissipation seems to self-adjust to the minimum required to prevent noise at the grid scale. This way of discretizing, in line with the ILES approach, turns out to be a viable alternative to conservative discretization combined with an explicit dissipation term.

At the bottom of Fig. 10 we plot the evolution of the total enstrophy at the three different resolutions. Once again, we note that the enstrophy dissipation decreases as the resolution increases: higher resolutions better preserve the PV variance, while low-resolution simulations suffer from more dissipation.

Table 1 Parameters of the idealized double-gyre configuration.

5.3 Double-gyre configuration

Our third numerical experiment to validate our solver is an idealized double-gyre configuration. Double-gyre configurations are a natural test for QG implementation or parameterization (e.g., Zanna et al., 2017; Ryzhov et al., 2020; Uchida et al., 2022). We consider here an octagonal ocean basin to illustrate the ability of our solver to handle non-square geometries. This octagon has maximal dimensions L_x×L_y. We assume free-slip boundary conditions on the boundaries. We consider N=3 layers on the vertical. We use an idealized stationary and symmetric wind stress (τ_x,τ_y) with ${\mathit{\tau }}_x=-({\mathit{\tau }}_{\mathrm{0}}/{\mathit{\rho }}_{\mathrm{0}})\cos (\mathrm{2}\mathit{\pi }y/L_y)$ and τ_y=0 on the top and linear drag at the bottom drag coefficient δ. The parameter values are given in Table 1.

Figure 11 (a, b) Upper-layer stream function ψ and relative vorticity (i.e., Δψ) after 50 years of spin-up. (c, d, e) Upper-layer mean stream function, mean kinetic energy density, and eddy kinetic energy density computed over 40 years after 10 years of spin-up.

We study this configuration in an eddy-permitting resolution of 20 km; the eddy-resolving meaning that the spatial resolution (20 km) is half of the larger baroclinic Rossby radius (41 km). Already at such eddy-permitting resolution, multi-layer QG solvers do not necessarily produce a well-pronounced eastward jet and usually require additional eddy parameterization (Uchida et al., 2022).

We plot at the top of Fig. 11 a snapshot of upper-layer SF and relative vorticity (i.e., Δψ) after 30 years of spin-up. As expected, our solver produces a strong western boundary current on the vertical and non-vertical boundaries. In the middle of this boundary starts a well-extended eastward jet whose length is qualitatively comparable with the Gulf Stream length in the North Atlantic basin. This jet is surrounded by a recirculation zone with several meso-scale eddies, which appears coherent with eddy-resolving-resolution simulations. We notice large-scale Rossby waves that are emerging near the eastern boundary and that propagate westward.

Figure 12 Upper-layer mean stream function, mean kinetic energy density, and eddy kinetic energy density computed over 40 years after 10 years of spin-up at resolutions (top to bottom) 27, 40, and 53 km.

Since this system is chaotic, showing a snapshot is not very representative of the dynamical behavior of the system. We plot statistics of our solution at the bottom of Fig. 11: the mean SF, the mean kinetic energy (i.e., the kinetic energy of the mean velocity), and the eddy kinetic energy (i.e., the kinetic energy of the velocity standard deviation). These statistics were computed over 40 years after 10 years of spin-up, saving one snapshot every 15 years. These statistics are symmetric, which is expected since the domain shape and the wind forcing are symmetric. They confirm the presence of a strong western boundary current and fluctuating eastward jet whose length is roughly three-quarters of the domain. These results seem to confirm the relevance of implicit dissipation provided by upwinding, since usual multi-layer QG solvers require additional eddy parameterization (Uchida et al., 2022) to produce a well-pronounced eastward jet.

To assess the influence of the resolution on the solution produced by our solver, we run the same double-gyre experiment at lower resolutions: 27, 40, and 53 km. We plot the mean SF as well as the mean and eddy kinetic energy statistics in Fig. 12. We note that the eastward jet progressively diminishes as the resolution decreases: at resolution 27 km, it barely reaches the middle of the domain, while at resolutions 40 and 53 km, it almost disappears. At these two coarsest resolutions that are comparable to the largest baroclinic Rossby radius, 41 km in our configuration, meso-scale eddies cannot be resolved properly by our solver, and one would require an additional eddy parameterization to produce the eastward jet (Zanna et al., 2017).

6 Conclusions

We presented MQGeometry, a multi-layer quasi-geostrophic equation solver for non-rectangular geometries. This solver has three original aspects compared to usual solvers such as Q-GCM, PyQG, or PEQUOD: the use of finite volume for advection via staggering of the PV and SF, the non-linear WENO upwind-biased reconstructions with implicit dissipation, and the ability to handle non-square geometry with a fast spectral DST solver combined with the capacitance matrix method. Running a simulation with this solver does not require the tuning of any additional parameter, e.g., additional hyper-viscosity.

This multi-layer QG solver delivers a computationally efficient playground for studying meso-scale non-linear dynamics. It opens the way to study QG dynamics in basins with a realistic coastline, e.g., Mediterranean or North Atlantic basins. Moreover, with PyTorch automatic differentiation, one can easily build upon this implementation to develop new machine-learning parameterizations of the QG sub-grid scales or new data-assimilation techniques using QG.

We believe that more complex modeling systems can be implemented in high-level languages like Python without sacrificing performance, as demonstrated by Häfner et al. (2021). The QG system that we implemented in this solver is fairly simple, and the present solver should be seen as a proof of concept. Our plan is to extend the presented approach to shallow-water equations and subsequently to primitive equations. Major advantages are the seamless parallelism offered by GPUs, enabling us to write code that closely aligns with the continuous equations, and the automatic differentiation, allowing us to learn vertical parameterization in an end-to-end fashion (Kochkov et al., 2023).

Appendix A: Usual discretization of multi-layer QG equations

One typically solves multi-layer QG equations using the strategy that follows (e.g., Hogg et al., 2014).

1. Use an evenly spaced Arakawa-C grid with the PV and the SF discretized on the same location, namely, the cell vertices (see left panel of Fig. 1).

2. Solve the PV advection Eq. (1a) using the energy–enstrophy-conservative Arakawa–Lamb scheme (Arakawa and Lamb, 1981) in the interior domain (i.e., not on the boundaries).

3. Since the scheme is energy conserving, use an additional (hyper-)viscosity scheme to dissipate the energy fueled by the wind forcing.

4. Given the matrix A diagonalization

   $${\displaystyle }\mathbf{A}={\mathbf{C}}_{\mathrm{m}\mathrm{2}\mathrm{l}}\mathbf{\Lambda }{\mathbf{C}}_{\mathrm{l}\mathrm{2}\mathrm{m}},$$

   where the layer-to-mode matrix C_l2m is the inverse of the mode-to-layer matrix C_m2l, and Λ is a diagonal matrix containing A eigenvalues, perform the following layer-to-mode transform:

   $${\displaystyle }\stackrel{\mathrm{̃}}{\mathit{\psi }},\stackrel{\mathrm{̃}}{\mathit{q}}={\mathbf{C}}_{\mathrm{l}\mathrm{2}\mathrm{m}}\mathit{\psi },{\mathbf{C}}_{\mathrm{l}\mathrm{2}\mathrm{m}}\mathit{q}.$$

   With this transform, the elliptic Eq. (1b) becomes a stack of N two-dimensional Helmholtz equations

   $$\begin{array}{}\text{(A1)}& {\mathrm{\Delta }}_{\mathrm{h}}\stackrel{\mathrm{̃}}{\mathit{\psi }}-f_{\mathrm{0}}^{\mathrm{2}}\mathbf{\Lambda }\stackrel{\mathrm{̃}}{\mathit{\psi }}=\stackrel{\mathrm{̃}}{\mathit{q}}-\mathit{\beta }y\end{array}$$

   with homogeneous Dirichlet boundary conditions.

5. Solve these N two-dimensional Helmholtz equations, using, e.g., fast diagonalization with type-I discrete sine transform (DST-I):

   $$\text{DST-I}[x{]}_k=\sum _{l=\mathrm{1}}^L{x}_l\sin \left[{\displaystyle \frac{\mathit{\pi }lk}{L+\mathrm{1}}}\right],k=\mathrm{1},\mathrm{\dots },L.$$

6. Transform back from mode to layers

   $${\displaystyle }\mathit{\psi },\mathit{q}={\mathbf{C}}_{\mathrm{m}\mathrm{2}\mathrm{l}}\stackrel{\mathrm{̃}}{\mathit{\psi }},{\mathbf{C}}_{\mathrm{m}\mathrm{2}\mathrm{l}}\stackrel{\mathrm{̃}}{\mathit{q}}.$$

7. Update the PV boundary values using the elliptic Eq. (1b). This requires definition of the Laplacian on the boundaries and possibly involves partial free-slip/no-slip boundary conditions.

Appendix B: WENO implementation

Listing B1 Python implementation of WENO-JS and WENO-Z reconstructions.

Appendix C: Vortex-shear instability

Figure C1 Vortex-shear instability final state (t=30τ) with resolution 512² and WENO-5 and with resolution 1024² and WENO-3/WENO-5.