jaxparrow

jaxparrow implements a novel approach based on a variational formulation to compute the inversion of the cyclogeostrophic balance.

It leverages the power of JAX, to efficiently solve the inversion as an optimization problem. Given the Sea Surface Height (SSH) field of an ocean system, jaxparrow estimates the velocity field that best satisfies the cyclogeostrophic balance.

See the full documentation!

Installation

jaxparrow is Pip-installable:

pip install jaxparrow

However, users with access to GPUs or TPUs should first install JAX separately in order to fully benefit from its high-performance computing capacities. See JAX instructions.
By default, jaxparrow will install a CPU-only version of JAX if no other version is already present in the Python environment.

Usage

As a package

The function you are most probably looking for is cyclogeostrophy. It computes the cyclogeostrophic velocity field (returned as two 2darray) from:

• a SSH field (a 2darray),
• the latitude and longitude grids at the T points (two 2darray).

In a Python script, assuming that the input grids have already been initialised / imported, estimating the cyclogeostrophic velocities for a single timestamp would resort to:

from jaxparrow import cyclogeostrophy

u_cyclo_2d, v_cyclo_2d = cyclogeostrophy(ssh_2d, lat_2d, lon_2d)

Because jaxparrow uses C-grids the velocity fields are represented on two grids (U and V), and the tracer fields (such as SSH) on one grid (T).
We provide functions computing some kinematics (such as velocities magnitude, normalized relative vorticity, or kinematic energy) accounting for these gridding system:

from jaxparrow.tools.kinematics import magnitude

uv_cyclo_2d, v_cyclo_2d = magnitude(u_cyclo_2d, v_cyclo_2d, interpolate=True)

To vectorise the estimation of the cyclogeostrophy across a first time dimension, one aims to use jax.vmap.

import jax

vmap_cyclogeostrophy = jax.vmap(cyclogeostrophy, in_axes=(0, None, None))
u_cyclo_3d, v_cyclo_3d = vmap_cyclogeostrophy(ssh_3d, lat_2d, lon_2d)

By default, the cyclogeostrophy function relies on our variational method. Its method argument provides the ability to use an iterative method instead, either the one described by Penven et al., or the one by Ioannou et al.. Additional arguments also give a finer control over the three approaches hyperparameters.
See jaxparrow API documentation for more details.

Notebooks are available as step-by-step examples.

As an executable

jaxparrow is also available from the command line:

jaxparrow --conf_path conf.yml

The YAML configuration file conf.yml instruct where input netCDF files are locally stored, and how to retrieve variables and coordinates from them. It also provides the path of the output netCDF file. Optionally, it can specify which cyclogeostrophic approach should be applied and its hyperparameters.

An example configuration file detailing all the required and optional entries can be found here.

Contributing

Contributions of all sorts are welcomed! See CONTRIBUTING.md and CONDUCT.md to get started.

How to cite

If you use this software, please cite it as specified in CITATION.cff. Thank you!