https://github.com/meom-group/jaxparrow
A package for computing the inversion of the cyclogeostrophic balance based on a variational formulation approach, using JAX.
https://github.com/meom-group/jaxparrow
cyclogeostrophy eddy flow geostrophy jax swirl velocity
Last synced: 5 months ago
JSON representation
A package for computing the inversion of the cyclogeostrophic balance based on a variational formulation approach, using JAX.
- Host: GitHub
- URL: https://github.com/meom-group/jaxparrow
- Owner: meom-group
- License: apache-2.0
- Created: 2023-10-10T12:05:22.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2025-12-01T08:21:15.000Z (7 months ago)
- Last Synced: 2025-12-03T19:48:48.762Z (7 months ago)
- Topics: cyclogeostrophy, eddy, flow, geostrophy, jax, swirl, velocity
- Language: Python
- Homepage: https://jaxparrow.readthedocs.io/
- Size: 118 MB
- Stars: 9
- Watchers: 3
- Forks: 1
- Open Issues: 4
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Citation: CITATION.cff
Awesome Lists containing this project
README
# jaxparrow
![Python](https://img.shields.io/badge/dynamic/yaml?url=https://raw.githubusercontent.com/meom-group/jaxparrow/master/.github/workflows/python-package.yml&label=Python&query=$.jobs.build.strategy.matrix["python-version"])
[](https://pypi.org/project/jaxparrow/)
[](https://jaxparrow.readthedocs.io/)
[](https://zenodo.org/badge/latestdoi/702998298)
`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`](https://jax.readthedocs.io/en/latest/), 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/](https://jaxparrow.readthedocs.io/en/latest/)!
## Installation
`jaxparrow` is Pip-installable:
```shell
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](https://jax.readthedocs.io/en/latest/installation.html).
By default, `jaxparrow` will install a CPU-only version of JAX if no other version is already present in the Python environment.
## Usage
The function you are most probably looking for is [`cyclogeostrophy`](https://jaxparrow.readthedocs.io/en/latest/api/#jaxparrow.cyclogeostrophy.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:
```python
from jaxparrow import cyclogeostrophy
ucg, vcg = cyclogeostrophy(ssh_2d, lat_2d, lon_2d, return_grids=False)
```
*Because `jaxparrow` uses [C-grids](https://xgcm.readthedocs.io/en/latest/grids.html) 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:
```python
from jaxparrow.tools.kinematics import magnitude
uv_cg = magnitude(ucg, vcg)
```
To vectorise the estimation of the cyclogeostrophy along a first time dimension, one aims to use `jax.vmap`.
```python
import jax
vmap_cyclogeostrophy = jax.vmap(cyclogeostrophy, in_axes=(0, None, None))
u_cg_3d, v_cg_3d = vmap_cyclogeostrophy(ssh_3d, lat_2d, lon_2d)
```
By default, the `cyclogeostrophy` function relies on our minimization-based method.
Its `method` argument provides the ability to use the fixed-point method instead, as described by [Penven *et al.* (2014)](https://doi.org/10.1016/j.dsr2.2013.10.015).
Additional arguments also give a finer control over the different approaches hyperparameters.
See `jaxparrow` [documentation](https://jaxparrow.readthedocs.io/en/latest/) for more details (including the API description and step-by-step examples).
## Contributing
Contributions are welcomed!
See [CONTRIBUTING.md](https://github.com/meom-group/jaxparrow/blob/main/CONTRIBUTING.md) and [CONDUCT.md](https://github.com/meom-group/jaxparrow/blob/main/CONDUCT.md) to get started.
## How to cite
If you use this software, please cite it: [CITATION.cff](https://github.com/meom-group/jaxparrow/blob/main/CITATION.cff).
Thank you!