jaxparrow

Description

Technical info

Usage

Estimating the cyclogeostrophic currents 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 and geometry modules.

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