GeophysicalFlows.jl: Solvers for geophysical fluid dynamics problems in periodic domains on CPUs & GPUs

GeophysicalFlows.jl is a Julia (Bezanson et al., 2017) package that contains partial differential equations solvers for a collection of geophysical fluid systems in periodic domains. All modules use Fourier-based pseudospectral numerical methods and leverage the framework provided by the FourierFlows.jl (Constantinou et al., 2021) Julia package for time-stepping, custom diagnostics, and saving output.


Statement of need
Conceptual models in simple domains often provide stepping stones for better understanding geophysical and astrophysical systems, particularly the atmospheres and oceans of Earth and other planets. These conceptual models are used in research but also are of great value for helping students in class to grasp new concepts and phenomena. Oftentimes people end up coding their own versions of solvers for the same partial differential equations for research or classwork. GeophysicalFlows.jl package is designed to be easily utilized and adaptable for a wide variety of both research and pedagogical purposes.
On top of the above-mentioned needs, the recent explosion of machine-learning applications in atmospheric and oceanic sciences advocates for the need that solvers for partial differential equations can be run on GPUs.
GeophysicalFlows.jl provides a collection of modules for solving sets of partial differential equations often used as conceptual models. These modules are continuously tested (unit tests and tests for the physics involved) and are well-documented. GeophysicalFlows.jl utilizes Julia's functionality and abstraction to enable all modules to run on CPUs or GPUs, and to provide a high level of customizability within modules. The abstractions allow simulations to be tailored for specific research questions, via the choice of parameters, domain properties, and schemes for damping, forcing, time-stepping etc. Simulations can easily be carried out on different computing architectures. Selection of the architecture on which equations are solved is done by providing the argument CPU() or GPU() during the construction of a particular problem.
Documented examples for each geophysical system (module) appear in the package's documentation, providing a starting point for new users and for the development of new or customized modules. Current modules include two-dimensional flow and a variety of quasigeostrophic (QG) dynamical systems, which provide analogues to the large-scale dynamics of atmospheres and oceans. The QG systems currently in GeophysicalFlows.jl extend two-dimensional dynamics to include the leading order effects of a third dimension through planetary rotation, topography, surface boundary conditions, stratification and quasi-twodimensional layering. A community-based collection of diagnostics throughout the modules are used to compute quantities like energy, enstrophy, dissipation, etc.
Figure 1: Potential vorticity snapshots from a nonlinearly equilibrated simulation of the Eady instability over a meridional ridge. Simulation used MultiLayerQG module of GeophysicalFlows.jl. The Eady problem was approximated here using 5 fluid layers stacked up in the vertical. Each layer was simulated with 512² grid-points. Plots were made with the Plots.jl Julia package, which utilizes the cmocean colormaps collection (Thyng et al., 2016). Scripts to reproduce the simulation reside in the repository github.com/FourierFlows/MultilayerQG-example.

State of the field
GeophysicalFlows.jl is a unique Julia package that shares some features and similarities with other packages. In particular: • pyqg (Abernathey et al., 2019) (Python) Beyond their base language, the major differences between GeophysicalFlows.jl and pyqg is that GeophysicalFlows.jl can be run on GPUs or CPUs and leverages a separate package (FourierFlows.jl; which is continuously developed) to solve differential equations and compute diagnostics, while pyqg can only be run on CPUs and uses a self-contained kernel.
• Dedalus (Burns et al., 2020) Dedalus is a Python package with an intuitive script-based interface that uses spectral methods to solve general partial differential equations, such as the ones within Geoph ysicalFlows.jl. Dedalus allows for more general boundary conditions in one of the dimensions. It only runs on CPUs (not on GPUs) but can be MPI-parallelized.

(Julia)
Oceananigans.jl is a fluid solver focussed on the Navier-Stokes equations under the Boussinesq approximation. Oceananigans.jl also runs on GPUs, and it allows for more variety of boundary conditions but it does not have spectral accuracy as it uses finite-volume discretization methods.
• MAOOAM (De Cruz et al., 2016) (Fortran, Python, and Lua) and its expanded Python implementation qgs (Demaeyer et al., 2020) MAOOAM and qgs simulate two atmospheric layers with QG dynamics, above either land or an oceanic fluid layer with reduced-gravity QG dynamics. The dynamics of individual layers have overlap with the MultiLayerQG and SingleLayerQG modules, however the layer configuration of MOAAM and qgs is specifically designed to study the dynamics of Earth's mid-latitude atmosphere. Neither MAOOAM nor qgs can run on GPUs.
• Isolated codes/scripts Several codes/scripts exist in personal websites and in open-source public repositories with similar functionality as some GeophysicalFlows.jl modules (e.g., TwoDNavier Stokes or SingleLayerQG). Usually, though, these codes come without any or poor documentation and typically they are not continuously tested.
GeophysicalFlows.jl can be used to investigate a variety of scientific research questions thanks to its various modules and high customizability, and its ease-of-use makes it an ideal teaching tool for fluids courses (Constantinou, 2020;Constantinou & Wagner, 2020). Geop hysicalFlows.jl has been used in developing Lagrangian vortices identification algorithms (Karrasch & Schilling, 2020) and to test new theories for diagnosing turbulent energy transfers in geophysical flows (Pearson et al., 2021). Currently, GeophysicalFlows.jl is being used, e.g., (i) to compare different observational sampling techniques in these flows, (ii) to study the bifurcation properties of Kolmogorov flows (Constantinou & Drivas, 2020), (iii) to study the genesis and persistence of the polygons of vortices present at Jovian high latitudes (Siegelman, Young, and Ingersoll; in prep), and (iv) to study how mesoscale macroturbulence affects mixing of tracers (Bisits & Constantinou, 2021).