jaxparrow package

Subpackages

Submodules

Module contents

References

jaxparrow.cyclogeostrophy(ssh_t: Float[Array, 'lat lon'], lat_t: Float[Array, 'lat lon'], lon_t: Float[Array, 'lat lon'], mask: Float[Array, 'lat lon'] | None = None, method: Literal['variational', 'iterative'] = 'variational', n_it: int | None = None, optim: GradientTransformation | str = 'sgd', optim_kwargs: dict | None = None, res_eps: float = 0.01, use_res_filter: bool = False, res_filter_size: int = 3, return_geos: bool = False, return_grids: bool = True, return_losses: bool = False) → [<class 'jaxtyping.Float[Array, 'lat lon']'>, Ellipsis]

Computes the cyclogeostrophic Sea Surface Current (SSC) velocity field from a Sea Surface Height (SSH) field using a variational (default) or iterative method.

The cyclogeostrophic SSC velocity field is computed on a C-grid, following NEMO convention [1].

Parameters:

  • ssh_t (Float[Array, "lat lon"]) – SSH field (on the T grid)

  • lat_t (Float[Array, "lat lon"]) – Latitude of the T grid

  • lon_t (Float[Array, "lat lon"]) – Longitude of the T grid

  • mask (Float[Array, "lat lon"], optional) –

    Mask defining the marine area of the spatial domain; 1 or True stands for masked (i.e. land)

    If not provided, inferred from ssh_t nan values

  • method (Literal["variational", "iterative"], optional) –

    Estimation method to use. If method="variational", then use our variational formulation. If method="iterative", then use an iterative approach [2], [3].

    Defaults to variational

  • n_it (int, optional) –

    Maximum number of iterations.

    Defaults to N_IT_VAR or N_IT_IT, based on method

  • optim (Union[optax.GradientTransformation, str], optional) –

    Optimizer to use. Can be an optax.GradientTransformation optimizer, or a string referring to such an optimizer [4].

    Defaults to sgd

  • optim_kwargs (dict, optional) –

    Optimizer arguments (such as learning rate, etc…).

    If None, only the learning rate is enforced to LR_VAR

  • res_eps (float, optional) –

    Residual tolerance of the iterative approach. When residuals are smaller, the iterative approach considers local convergence to cyclogeostrophy.

    Defaults to RES_EPS_IT

  • use_res_filter (bool, optional) –

    Use of a convolution filter for the iterative approach when computing the residuals [3] or not [2].

    Defaults to False

  • res_filter_size (int, optional) –

    Size of the convolution filter of the iterative approach, when use_res_filter=True.

    Defaults to RES_FILTER_SIZE_IT

  • return_geos (bool, optional) –

    If True, returns the geostrophic SSC velocity field in addition to the cyclogeostrophic one.

    Defaults to False

  • return_grids (bool, optional) –

    If True, returns the U and V grids.

    Defaults to True

  • return_losses (bool, optional) –

    If True, returns the losses (cyclogeostrophic imbalance) over iterations.

    Defaults to False

Returns:

  • u_cyclo_u (Float[Array, “lat lon”]) – U component of the cyclogeostrophic SSC velocity field (on the U grid)

  • v_cyclo_v (Float[Array, “lat lon”]) – V component of the cyclogeostrophic SSC velocity field (on the V grid)

  • lat_u (Float[Array, “lat lon”]) – Latitudes of the U grid, if return_grids=True

  • lon_u (Float[Array, “lat lon”]) – Longitudes of the U grid, if return_grids=True

  • lat_v (Float[Array, “lat lon”]) – Latitudes of the V grid, if return_grids=True

  • lon_v (Float[Array, “lat lon”]) – Longitudes of the V grid, if return_grids=True

  • u_geos_u (Float[Array, “lat lon”]) – U component of the geostrophic SSC velocity field (on the U grid), if return_geos=True

  • v_geos_v (Float[Array, “lat lon”]) – V component of the geostrophic SSC velocity field (on the V grid), if return_geos=True

  • losses (Float[Array, “n_it”]) – Cyclogeostrophic imbalance evaluated at each iteration, if return_losses=True

jaxparrow.geostrophy(ssh_t: Float[Array, 'lat lon'], lat_t: Float[Array, 'lat lon'], lon_t: Float[Array, 'lat lon'], mask: Float[Array, 'lat lon'] | None = None, return_grids: bool = True) → [<class 'jaxtyping.Float[Array, 'lat lon']'>, Ellipsis]

Computes the geostrophic Sea Surface Current (SSC) velocity field from a Sea Surface Height (SSH) field.

The geostrophic SSC velocity field is computed on a C-grid, following NEMO convention [1].

Parameters:

  • ssh_t (Float[Array, "lat lon"]) – SSH field (on the T grid)

  • lat_t (Float[Array, "lat lon"]) – Latitudes of the T grid

  • lon_t (Float[Array, "lat lon"]) – Longitudes of the T grid

  • mask (Float[Array, "lat lon"], optional) –

    Mask defining the marine area of the spatial domain; 1 or True stands for masked (i.e. land)

    If not provided, inferred from ssh_t nan values

  • return_grids (bool, optional) –

    If True, returns the U and V grids.

    Defaults to True

Returns:

  • u_geos_u (Float[Array, “lat lon”]) – U component of the geostrophic SSC velocity field (on the U grid)

  • v_geos_v (Float[Array, “lat lon”]) – V component of the geostrophic SSC velocity field (on the V grid)

  • lat_u (Float[Array, “lat lon”]) – Latitudes of the U grid, if return_grids=True

  • lon_u (Float[Array, “lat lon”]) – Longitudes of the U grid, if return_grids=True

  • lat_v (Float[Array, “lat lon”]) – Latitudes of the V grid, if return_grids=True

  • lon_v (Float[Array, “lat lon”]) – Longitudes of the V grid, if return_grids=True