Computes the inversion of the cyclogeostrophic balance based on a variational formulation approach, using JAX
Project description
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) or the geostrophic velocity field of an ocean system, jaxparrow estimates the velocity field that best satisfies the cyclogeostrophic balance.
Installation
The package 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
Two functions are directly available from jaxparrow
:
geostrophy
: computes the geostrophic velocity field (returns twonumpy 2darray
) from a SSH2darray
, two2darray
of spatial steps, and two2darray
of Coriolis factors.cyclogeostrophy
: computes the cyclogeostrophic velocity field (returns two2darray
) from two2darray
of geostrophic velocities, four2darray
of spatial steps, and two2darray
of Coriolis factors.
Because jaxparrow uses C-grids the velocity fields are represented on two grids, and the SSH on one grid.
In a Python script, assuming that the input grids have already been initialised / imported, it would simply resort to:
from jaxparrow import cyclogeostrophy, geostrophy
u_geos, v_geos = geostrophy(ssh=ssh,
dx_ssh=dx_ssh, dy_ssh=dy_ssh,
coriolis_factor_u=coriolis_factor_u, coriolis_factor_v=coriolis_factor_v)
u_cyclo, v_cyclo = cyclogeostrophy(u_geos=u_geos, v_geos=v_geos,
dx_u=dx_u, dx_v=dx_v, dy_u=dy_u, dy_v=dy_v,
coriolis_factor_u=coriolis_factor_u, coriolis_factor_v=coriolis_factor_v)
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.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for jaxparrow-0.0.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 43cea37b3cf1353def73a8b768e6d953e8728ee280211119bfcaddd97aefcd55 |
|
MD5 | 0c892b92f738ad08b75340af823d8fac |
|
BLAKE2b-256 | ca7c476251135dcb7e0f79ec60c72279fd937c3e81a5ab4dc55adf1be1e5a56f |