# 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/)
***jaxparrow*** implements a novel approach based on a variational 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 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.
## Installation
The package is Pip-installable:
```shell
pip install jaxparrow
```
**<ins>However</ins>**, 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
### As a package
Two functions are directly available from `jaxparrow`:
- `geostrophy` computes the geostrophic velocity field (returns two `2darray`) from:
- a SSH field (a `2darray`),
- the latitude and longitude at the T points (two `2darray`),
- an optional mask grid (one `2darray`).
- `cyclogeostrophy` computes the cyclogeostrophic velocity field (returns two `2darray`) from:
- a SSH field (a `2darray`),
- the latitude and longitude at the T points (two `2darray`),
- an optional mask grid (one `2darray`).
*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 SSH on one grid (T).*
In a Python script, assuming that the input grids have already been initialised / imported, it would resort to:
```python
from jaxparrow import cyclogeostrophy, geostrophy
u_geos, v_geos = geostrophy(ssh_t=ssh,
lat_t=lat, lon_t=lon,
mask=mask)
u_cyclo, v_cyclo = cyclogeostrophy(ssh_t=ssh,
lat_t=lat, lon_t=lon,
mask=mask)
```
To vectorise the application of the `geostrophy` and `cyclogeostrophy` functions across an added time dimension, one aims to utilize `vmap`.
However, this necessitates avoiding the use of `np.ma.masked_array`.
Hence, our functions accommodate mask `array` as parameter to effectively consider masked regions.
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.*](https://doi.org/10.1016/j.dsr2.2013.10.015), or the one by [Ioannou *et al.*](https://doi.org/10.1029/2019JC015031).
Additional arguments also give a finer control over the three approaches hyperparameters. \
See **jaxparrow** [API documentation](https://jaxparrow.readthedocs.io/en/latest/api.html) for more details.
[Notebooks](https://jaxparrow.readthedocs.io/en/latest/examples.html) are available as step-by-step examples.
### As an executable
**jaxparrow** is also available from the command line:
```shell
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](https://github.com/meom-group/jaxparrow/blob/main/docs/example-conf.yml).
Raw data
{
"_id": null,
"home_page": null,
"name": "jaxparrow",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "cyclogeostrophy, eddy, flow, geostrophy, jax, swirl, velocity",
"author": null,
"author_email": "Vadim Bertrand <vadim.bertrand@univ-grenoble-alpes.fr>, Victor E V Z De Almeida <victorzaia@outlook.com>, Julien Le Sommer <julien.lesommer@univ-grenoble-alpes.fr>, Emmanuel Cosme <emmanuel.cosme@univ-grenoble-alpes.fr>",
"download_url": "https://files.pythonhosted.org/packages/a9/bf/960fc85cece1c7b6987b18b44a92a001306aa78ec4c3ce63167ace9554fa/jaxparrow-0.2.2.tar.gz",
"platform": null,
"description": "# jaxparrow\n\n![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\"])\n[](https://pypi.org/project/jaxparrow/)\n\n[](https://jaxparrow.readthedocs.io/)\n\n***jaxparrow*** implements a novel approach based on a variational formulation to compute the inversion of the cyclogeostrophic balance.\n\nIt leverages the power of [JAX](https://jax.readthedocs.io/en/latest/), to efficiently solve the inversion as an optimization problem. \nGiven the Sea Surface Height (SSH) field of an ocean system, **jaxparrow** estimates the velocity field that best satisfies the cyclogeostrophic balance.\n\n## Installation\n\nThe package is Pip-installable:\n```shell\npip install jaxparrow\n```\n\n**<ins>However</ins>**, users with access to GPUs or TPUs should first install JAX separately in order to fully benefit from its high-performance computing capacities. \nSee [JAX instructions](https://jax.readthedocs.io/en/latest/installation.html). \\\nBy default, **jaxparrow** will install a CPU-only version of JAX if no other version is already present in the Python environment.\n\n## Usage\n\n### As a package\n\nTwo functions are directly available from `jaxparrow`:\n\n- `geostrophy` computes the geostrophic velocity field (returns two `2darray`) from:\n - a SSH field (a `2darray`), \n - the latitude and longitude at the T points (two `2darray`), \n - an optional mask grid (one `2darray`).\n- `cyclogeostrophy` computes the cyclogeostrophic velocity field (returns two `2darray`) from:\n - a SSH field (a `2darray`), \n - the latitude and longitude at the T points (two `2darray`), \n - an optional mask grid (one `2darray`).\n\n*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 SSH on one grid (T).*\n\nIn a Python script, assuming that the input grids have already been initialised / imported, it would resort to:\n\n```python\nfrom jaxparrow import cyclogeostrophy, geostrophy\n\nu_geos, v_geos = geostrophy(ssh_t=ssh,\n lat_t=lat, lon_t=lon,\n mask=mask)\nu_cyclo, v_cyclo = cyclogeostrophy(ssh_t=ssh,\n lat_t=lat, lon_t=lon,\n mask=mask)\n```\n\nTo vectorise the application of the `geostrophy` and `cyclogeostrophy` functions across an added time dimension, one aims to utilize `vmap`.\nHowever, this necessitates avoiding the use of `np.ma.masked_array`. \nHence, our functions accommodate mask `array` as parameter to effectively consider masked regions.\n\nBy default, the `cyclogeostrophy` function relies on our variational method.\nIts `method` argument provides the ability to use an iterative method instead, either the one described by [Penven *et al.*](https://doi.org/10.1016/j.dsr2.2013.10.015), or the one by [Ioannou *et al.*](https://doi.org/10.1029/2019JC015031).\nAdditional arguments also give a finer control over the three approaches hyperparameters. \\\nSee **jaxparrow** [API documentation](https://jaxparrow.readthedocs.io/en/latest/api.html) for more details.\n\n[Notebooks](https://jaxparrow.readthedocs.io/en/latest/examples.html) are available as step-by-step examples.\n\n### As an executable\n\n**jaxparrow** is also available from the command line:\n```shell\njaxparrow --conf_path conf.yml\n```\nThe YAML configuration file `conf.yml` instruct where input netCDF files are locally stored, and how to retrieve variables and coordinates from them.\nIt also provides the path of the output netCDF file. Optionally, it can specify which cyclogeostrophic approach should be applied and its hyperparameters.\n\nAn example configuration file detailing all the required and optional entries can be found [here](https://github.com/meom-group/jaxparrow/blob/main/docs/example-conf.yml).\n",
"bugtrack_url": null,
"license": "Apache-2.0",
"summary": "Computes the inversion of the cyclogeostrophic balance based on a variational formulation approach, using JAX",
"version": "0.2.2",
"project_urls": {
"Bug Tracker": "https://github.com/meom-group/jaxparrow/issues",
"Homepage": "https://jaxparrow.readthedocs.io/"
},
"split_keywords": [
"cyclogeostrophy",
" eddy",
" flow",
" geostrophy",
" jax",
" swirl",
" velocity"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "6cfff465e6a46b61cb014f69507428371acb213ffc6d9fdaa7225ac12ebff6fd",
"md5": "849f8b982e92aa475ba17d8c03c449ca",
"sha256": "e43ae8ce34acecadca3a1e61f96332a9cfb5fc2e468caf315e1b21d273bb6fe5"
},
"downloads": -1,
"filename": "jaxparrow-0.2.2-py3-none-any.whl",
"has_sig": false,
"md5_digest": "849f8b982e92aa475ba17d8c03c449ca",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 25792,
"upload_time": "2024-08-23T15:39:07",
"upload_time_iso_8601": "2024-08-23T15:39:07.977933Z",
"url": "https://files.pythonhosted.org/packages/6c/ff/f465e6a46b61cb014f69507428371acb213ffc6d9fdaa7225ac12ebff6fd/jaxparrow-0.2.2-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "a9bf960fc85cece1c7b6987b18b44a92a001306aa78ec4c3ce63167ace9554fa",
"md5": "57280fe57c0d2d68b8a69d8897d6b0ee",
"sha256": "f5cfcb76de7a7b08011325ae46612d0d0e8a63dcc68e85f66773a0a99310cffb"
},
"downloads": -1,
"filename": "jaxparrow-0.2.2.tar.gz",
"has_sig": false,
"md5_digest": "57280fe57c0d2d68b8a69d8897d6b0ee",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 23022,
"upload_time": "2024-08-23T15:39:09",
"upload_time_iso_8601": "2024-08-23T15:39:09.520817Z",
"url": "https://files.pythonhosted.org/packages/a9/bf/960fc85cece1c7b6987b18b44a92a001306aa78ec4c3ce63167ace9554fa/jaxparrow-0.2.2.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-08-23 15:39:09",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "meom-group",
"github_project": "jaxparrow",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "jaxparrow"
}
JSON
