jaxparrow

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

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

A comprehensive documenation is available: https://jaxparrow.readthedocs.io/en/latest/!

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

Estimating the cyclogeostrophic currents from a given Sea Surface Height field can be achieved using any of the following methods:

• minimization_based,
• gradient_wind,
• fixed_point.

Taking as inputs:

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

They return a result objects holding the cyclogeostrophic velocity $u$ and $v$ components.

In a Python script estimating the cyclogeostrophic currents for a single timestamp would resort to:

from jaxparrow import minimization_based  # or gradient_wind or fixed_point

mb_result = minimization_based(lat_t=lat_2d, lon_t=lon_2d, ssh_t=ssh_2d)

ucg_2d = mb_result.ucg  # 2d jax.Array
vcg_2d = mb_result.vcg  # 2d jax.Array

Note that it is also possible to directly pass as inputs the geostrophic velocity $u$ and $v$ components, rather than the SSH:

mb_result = minimization_based(lat_t=lat_2d, lon_t=lon_2d, ug_t=ug_2d, vg_t=vg_2d)

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

import jax

vmap_cyclogeostrophy = jax.vmap(lambda _ssh_2d: cyclogeostrophy(lat_t=lat_2d, lon_t=lon_2d, ssh_t=_ssh_2d))
mb_result = vmap_cyclogeostrophy(ssh_3d)

ucg_3d = mb_result.ucg  # 3d jax.Array
vcg_3d = mb_result.vcg  # 3d jax.Array

jaxparrow also notably allows to:

• add arbitrary regularization term when employing the minimization-based method, see the regularization with SWOT data example notebook,
• reconstruct geostrophic currents, see the geostrophy function,
• compute some classical kinematics, see the kinematics module,
• perform standard operations on the grid, see the operators module.

Explore jaxparrow documentation for more details, including the API description and step-by-step examples in the form of notebooks.

Contributing

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

How to cite

If you use this software, please cite it: CITATION.cff. Thank you!