GCM adapters#

Synthetic (theoretical) GCM interface for TRACMASS JAX.

Provides functions that generate idealised, analytically-defined velocity flux fields and initial particle states for testing and validation of the TRACMASS particle-tracking framework without requiring real ocean reanalysis data.

tracmass.gcm.theoretical.generate_keys(random_seed=42)[source]#

Generate a set of JAX random keys from a single seed.

Splits a PRNGKey derived from random_seed into 10 independent sub-keys for use in stochastic field generation.

Parameters:: random_seed (int, optional) – Integer seed used to initialise the JAX PRNG. Default is 42.
Returns:: jnp.ndarray – Array of shape (10, 2) containing 10 independent JAX random keys.

tracmass.gcm.theoretical.fluxes(dtm='2007-06-01', imt=100, jmt=80, km=30, random_seed=42)[source]#

Generate synthetic volume-flux fields on an idealised rectangular grid.

Constructs analytically-defined zonal and meridional volume fluxes following a double-gyre pattern with small random perturbations, and random vertical fluxes. When dtm is later than the reference date "2007-06-01", the fluxes are time-evolved by accumulating day-by-day random increments.

Backward-in-time tracking is handled at the new_grid() / new_flux() level by passing dtm_end < dtm_start; this function always returns forward (positive) fluxes.

Parameters:

dtm (pd.Timestamp or str, optional) – Target date for the flux snapshot. Default is "2007-06-01".
imt (int, optional) – Number of grid cells in the x-direction. Default is 100.
jmt (int, optional) – Number of grid cells in the y-direction. Default is 80.
km (int, optional) – Number of vertical levels. Default is 30.
random_seed (int, optional) – Seed for JAX random number generation. Default is 42.

Returns:

FluxDataset – xarray Dataset containing the variables required by GridFields: uflux, vflux, wflux, kmt, dzt, dxdy, timestep, dtm, reverse, imt, jmt, km, iperio, jperio, and zstot.

Return type:

FluxDataset

tracmass.gcm.theoretical.initial_state(n_particles=10_000, imt=100, jmt=80, km=30, random_seed=42)[source]#

Generate a random initial particle state on the theoretical grid.

Draws particle positions uniformly at random within the interior of the grid (avoiding boundary cells) and computes the corresponding integer cell indices.

Parameters:

n_particles (int, optional) – Number of particles to initialise. Default is 10_000.
imt (int, optional) – Number of grid cells in the x-direction. Default is 100.
jmt (int, optional) – Number of grid cells in the y-direction. Default is 80.
km (int, optional) – Number of vertical levels. Default is 30.
random_seed (int, optional) – Seed for JAX random number generation. Default is 42.

Returns:

ParticleState – Named tuple containing particle positions (x, y, z), integer cell indices (ia, ja, ka), time trackers (tt, ts, t0), sub-volume (subvol), iteration counter (niter), and active mask (active).

Return type:

ParticleState

Fabboni analytical test-case GCM for TRACMASS JAX.

Provides spatially-uniform, analytically-defined velocity fluxes and the canonical single-particle initial state matching the Fabboni test case from the TRACMASS Fortran distribution (projects/Theoretical).

The velocity field decays exponentially from initial conditions u = U0 = 0.5 m/s, v = 0, with inertial oscillation governed by the Coriolis parameter at 45° N and two damping time-scales GAMMA and GAMMAG. Because the field is spatially uniform, all particles experience identical forcing — making it ideal for comparison with the Fortran reference trajectory.

This module mirrors the interface of tracmass.gcm.theoretical:

fluxes() — returns an xr.Dataset compatible with new_flux() and update_flux().
initial_state() — returns a single-particle ParticleState that matches the Fortran seeding (isec=1, idir=1).

tracmass.gcm.fabboni.fluxes(dtm='2000-01-01')[source]#

Generate Fabboni spatially-uniform flux fields at time dtm.

The step index is derived from the offset of dtm from DTM0 = 2000-01-01, counted in units of TIMESTEP (1 hour).

Backward-in-time tracking is handled at the new_grid() / new_flux() level by passing dtm_end < dtm_start; this function always returns forward (positive) fluxes at the requested time.

Parameters:: dtm (pd.Timestamp or str, optional) – Target time for the flux snapshot. Default "2000-01-01".
Returns:: xr.Dataset – Dataset containing uflux, vflux, wflux, dzt, zstot, kmt, dxdy, timestep, dtm, reverse, imt, jmt, km, iperio, jperio.
Return type:: Dataset

tracmass.gcm.fabboni.initial_state()[source]#

Create the Fabboni reference initial particle state.

Seeds a single particle at the eastern wall of cell (150, 100, 1), matching the Fortran seeding isec=1, idir=1:

x = 150.0 — eastern wall of i-cell 150 → ia = 151
y = 99.5 — centre of j-cell 100 → ja = 100
z = 0.5 — centre of k-cell 1 → ka = 1

Returns:: ParticleState – Single-particle state with all time counters at zero and subvol set to one TIMESTEP in seconds.
Return type:: ParticleState

GLORYS12 reanalysis GCM interface for TRACMASS JAX.

This module provides a fluxes callable that loads GLORYS12 ocean reanalysis velocity fields and converts them to volume-flux arrays suitable for use with the TRACMASS particle-tracking framework.

tracmass.gcm.glorys.fluxes(dtm=None, include_wflux=True)[source]#

Load GLORYS12 flux fields for a single time window.

Reads GLORYS12 reanalysis velocities and grid geometry for the requested time, computes zonal, meridional, and vertical volume fluxes, and returns them as an xr.Dataset compatible with tracmass.gcm.fluxes_to_grid().

Backward-in-time tracking is handled at the new_grid() / new_flux() level by passing dtm_end < dtm_start; this function always returns forward (positive) fluxes.

Parameters:

dtm (pd.Timestamp or str or None, optional) – Target time for the flux snapshot. Default None — no dtm is passed to tracmass.seagrid.glorys.open_dataset() at all, so its own default applies.
include_wflux (bool)

Returns:

FluxDataset – xarray Dataset containing the variables required by GridFields: uflux, vflux, wflux, kmt, dzt, dxdy, timestep, dtm, reverse, imt, jmt, km, iperio, jperio, and zstot.

Return type:

FluxDataset

tracmass.gcm.glorys.lonlat_grid(grid=None)[source]#

Cell-centre longitude / latitude of a run’s grid.

Parameters:: grid (GridFields or None, optional) – The run grid; its i_offset / j_offset / imt / jmt select the matching window of the full GLORYS12 grid (a subset). Default: the full grid.
Returns:: lon, lat (np.ndarray) – Degrees, shape (imt, jmt) — indexed like grid.kmt.
Return type:: tuple[ndarray, ndarray]

tracmass.gcm.glorys.lonlat_to_xy(lon, lat)[source]#

Convert geographic coordinates to fractional TRACMASS (x, y) grid positions.

The exact inverse of xy_to_lonlat() (see tracmass.trajs.lonlat_to_xy()): a point off-centre within a cell gets the matching fractional offset rather than being snapped to the cell centre. A point in a land cell (kmt == 0) is placed at the centre of the nearest ocean cell instead.

The resulting (x, y) values can be passed directly to tracmass.seed.expand_sub() or used to initialise a ParticleState, with ia = floor(x) + 1, ja = floor(y) + 1 (1-based cell indices).

Parameters:

lon (array-like) – Longitude and latitude in degrees.
lat (array-like) – Longitude and latitude in degrees.

Returns:

x, y (np.ndarray) – Fractional 0-based TRACMASS grid positions, same shape as lon/lat.

Return type:

tuple[ndarray, ndarray]

Notes

The same x-halo trim as fluxes() is applied. Positions are full-grid (absolute) indices — the same frame as saved trajectory files — regardless of any subset. Subtract grid.i_offset / grid.j_offset to get positions relative to a subset grid (e.g. for seeding).

Examples

>>> import numpy as np
>>> from tracmass.gcm import glorys as gcm
>>> lon = np.array([-30.0, 10.0])
>>> lat = np.array([45.0, 55.0])
>>> x, y = gcm.lonlat_to_xy(lon, lat)

tracmass.gcm.glorys.xy_to_lonlat(x, y)[source]#

Convert fractional TRACMASS (x, y) positions to longitude/latitude.

Applies the same x-halo trim that fluxes() uses (but not the subset: x/y are full-grid indices, as saved by the trajectory recorder), then bilinearly interpolates the curvilinear grid’s lon/lat arrays at the fractional particle positions. Works correctly for curvilinear grids (e.g. NEMO ORCA) because interpolation is done in Cartesian (x, y, z) space on the unit sphere to avoid longitude discontinuities at ±180°.

Parameters:

x (array-like) – Fractional 0-based grid positions as stored in the trajectory Parquet file (columns "x" and "y"). floor(x) is the cell (matching column "ia" - 1 / "ja" - 1); its centre, where the grid’s lon/lat are defined, is at x = floor(x) + 0.5.
y (array-like) – Fractional 0-based grid positions as stored in the trajectory Parquet file (columns "x" and "y"). floor(x) is the cell (matching column "ia" - 1 / "ja" - 1); its centre, where the grid’s lon/lat are defined, is at x = floor(x) + 0.5.

Returns:

lon, lat (np.ndarray) – Longitude and latitude in degrees, same shape as x and y.

Return type:

tuple[ndarray, ndarray]

Notes

The grid-definition file caches longitude and latitude from the NEMO mesh_hgr file. If the file was built before the nav_lat bug-fix was applied, delete it so it is regenerated:

from tracmass.seagrid import glorys as _g
_g.grid_definition(force=True)

Examples

>>> import pandas as pd
>>> df = pd.read_parquet("trm_omi.parquet", columns=["x", "y"])
>>> lon, lat = gcm.glorys.xy_to_lonlat(df["x"], df["y"])

ECCO2 cube92 GCM interface for TRACMASS JAX.

This module provides a fluxes callable that loads ECCO2 cube92 daily ocean state estimate velocity fields and converts them to volume-flux arrays suitable for use with the TRACMASS particle-tracking framework.

Unlike GLORYS12/NEMO, ECCO2 cube92 is a plain global regular lat/lon grid (no curvilinear coordinates, no halo columns), so fluxes() needs no x-halo trimming and cell face widths (dxv, dyu) come from tracmass.seagrid.ecco2.grid_definition() rather than a mesh file.

tracmass.gcm.ecco2.fluxes(dtm=None, include_wflux=None)[source]#

Load ECCO2 cube92 flux fields for a single time window.

Reads ECCO2 cube92 velocities and grid geometry for the requested time, computes zonal, meridional, and (optionally) vertical volume fluxes, and returns them as an xr.Dataset compatible with tracmass.gcm.fluxes_to_grid().

Backward-in-time tracking is handled at the new_grid() / new_flux() level by passing dtm_end < dtm_start; this function always returns forward (positive) fluxes.

Parameters:

dtm (pd.Timestamp or str or None, optional) – Target time for the flux snapshot. Default None — no dtm is passed to tracmass.seagrid.ecco2.open_dataset() at all, so its own default applies.
include_wflux (bool or None, optional) – If True, also fetch wvel and compute wflux. If False, the wvel OPeNDAP dataset is not touched at all — use this with advection_scheme.piecewise(grid, no_wflux=True) for horizontal-only advection. If None (default), taken from the move_vertically setting (default True); see tracmass.config.move_vertically().

Returns:

FluxDataset – xarray Dataset containing the variables required by GridFields: uflux, vflux, (wflux), kmt, dzt, dxdy, timestep, dtm, reverse, imt, jmt, km, iperio, jperio, and zstot.

Return type:

FluxDataset

tracmass.gcm.ecco2.lonlat_grid(grid=None)[source]#

Cell-centre longitude / latitude of a run’s grid.

Parameters:: grid (GridFields or None, optional) – The run grid; its i_offset / j_offset / imt / jmt select the matching window of the full ECCO2 grid (a subset). Default: the full grid.
Returns:: lon, lat (np.ndarray) – Degrees, shape (imt, jmt) — indexed like grid.kmt.
Return type:: tuple[ndarray, ndarray]

tracmass.gcm.ecco2.lonlat_to_xy(lon, lat)[source]#

Convert geographic coordinates to fractional TRACMASS (x, y) grid positions.

The exact inverse of xy_to_lonlat() (see tracmass.trajs.lonlat_to_xy()): a point off-centre within a cell gets the matching fractional offset rather than being snapped to the cell centre. A point in a land cell (kmt == 0) is placed at the centre of the nearest ocean cell instead.

The resulting (x, y) values can be passed directly to tracmass.seed.expand_sub() or used to initialise a ParticleState, with ia = floor(x) + 1, ja = floor(y) + 1 (1-based cell indices).

Parameters:

lon (array-like) – Longitude and latitude in degrees.
lat (array-like) – Longitude and latitude in degrees.

Returns:

x, y (np.ndarray) – Fractional 0-based TRACMASS grid positions, same shape as lon/lat.

Return type:

tuple[ndarray, ndarray]

Notes

Positions are full-grid (absolute) indices — the same frame as saved trajectory files — regardless of any subset. Subtract grid.i_offset / grid.j_offset to get positions relative to a subset grid (e.g. for seeding).

Examples

>>> import numpy as np
>>> from tracmass.gcm import ecco2 as gcm
>>> lon = np.array([-30.0, 10.0])
>>> lat = np.array([45.0, 55.0])
>>> x, y = gcm.lonlat_to_xy(lon, lat)

tracmass.gcm.ecco2.xy_to_lonlat(x, y)[source]#

Convert fractional TRACMASS (x, y) positions to longitude/latitude.

x/y are full-grid indices (as saved by the trajectory recorder), whatever the subset. Bilinearly interpolates the grid’s lon/lat arrays at the fractional particle positions.

Parameters:

x (array-like) – Fractional 0-based grid positions as stored in the trajectory Parquet file (columns "x" and "y"). floor(x) is the cell (matching column "ia" - 1 / "ja" - 1); its centre, where the grid’s lon/lat are defined, is at x = floor(x) + 0.5.
y (array-like) – Fractional 0-based grid positions as stored in the trajectory Parquet file (columns "x" and "y"). floor(x) is the cell (matching column "ia" - 1 / "ja" - 1); its centre, where the grid’s lon/lat are defined, is at x = floor(x) + 0.5.

Returns:

lon, lat (np.ndarray) – Longitude and latitude in degrees, same shape as x and y.

Return type:

tuple[ndarray, ndarray]

Examples

>>> import pandas as pd
>>> df = pd.read_parquet("trm_ecco2.parquet", columns=["x", "y"])
>>> lon, lat = gcm.ecco2.xy_to_lonlat(df["x"], df["y"])

Scripps Gulf of Guinea (GoG) MITgcm interface for TRACMASS JAX.

This module provides a fluxes callable that loads the daily Scripps GoG regional MITgcm velocities (top 200 m) and converts them to volume-flux arrays for the TRACMASS particle-tracking framework.

The files keep MITgcm’s C-grid staggering: UVEL[i] sits on the west face of T-cell i and VVEL[j] on the south face of T-cell row j (with one extra row for the northernmost north face). TRACMASS reads uflux[i] / vflux[:, j] as the east / north faces, so fluxes() shifts both by one face. The east face of the last column is not in the files and is set to zero (it is land).

Transports use the files’ own face metrics and open fractions, uflux = UVEL * dyG * drF * hFacW and vflux = VVEL * dxG * drF * hFacS. The western (11 W) and southern (2 N) open boundaries carry inflow in the files, but TRACMASS treats a grid edge as closed, so those faces are zero in the tracking (particles cannot leave or enter there). The files hold no vertical velocity, so tracking is horizontal-only.

tracmass.gcm.scrippsgog.face_transports(vels)[source]#

Volume transports through the east and north face of every T-cell.

Parameters:: vels (xarray.Dataset) – A velocity file from tracmass.seagrid.scrippsgog.open_dataset() (staggered (z, YC, XG) / (z, YG, XC) dims).
Returns:: uflux, vflux (xarray.DataArray) – m³/s on T-cell dims (z, y, x). uflux[..., i] is the east face of cell i (zero for the last column); vflux[:, j] is the north face of row j.
Return type:: tuple[DataArray, DataArray]

tracmass.gcm.scrippsgog.fluxes(dtm=None)[source]#

Load Scripps GoG flux fields for a single time window.

Reads the daily velocity file for the requested time (downloading it if needed), computes zonal and meridional volume fluxes on the TRACMASS east/north face convention, and returns them as an xr.Dataset compatible with tracmass.gcm.fluxes_to_grid().

Backward-in-time tracking is handled at the new_grid() / new_flux() level by passing dtm_end < dtm_start; this function always returns forward (positive) fluxes.

Parameters:: dtm (pd.Timestamp or str or None, optional) – Target day. Default None — no dtm is passed to tracmass.seagrid.scrippsgog.open_dataset(), so its own default applies.
Returns:: FluxDataset – xarray Dataset with uflux, vflux, kmt, dzt, dxdy, timestep, dtm, reverse, imt, jmt, km, iperio, jperio and zstot. There is no wflux (the files have no vertical velocity).
Return type:: FluxDataset

tracmass.gcm.scrippsgog.lonlat_grid(grid=None)[source]#

Cell-centre longitude / latitude of a run’s grid.

Parameters:: grid (GridFields or None, optional) – The run grid; its i_offset / j_offset / imt / jmt select the matching window of the full GoG grid (a subset). Default: the full grid.
Returns:: lon, lat (np.ndarray) – Degrees, shape (imt, jmt) — indexed like grid.kmt.
Return type:: tuple[ndarray, ndarray]

tracmass.gcm.scrippsgog.lonlat_to_xy(lon, lat)[source]#

Convert geographic coordinates to fractional TRACMASS (x, y) grid positions.

The exact inverse of xy_to_lonlat() (see tracmass.trajs.lonlat_to_xy()). A point in a land cell (kmt == 0) is placed at the centre of the nearest ocean cell.

Parameters:

lon (array-like) – Longitude and latitude in degrees.
lat (array-like) – Longitude and latitude in degrees.

Returns:

x, y (np.ndarray) – Fractional 0-based full-grid (absolute) TRACMASS positions, same shape as lon/lat. Subtract grid.i_offset / grid.j_offset for positions relative to a subset grid.

Return type:

tuple[ndarray, ndarray]

Examples

>>> from tracmass.gcm import scrippsgog as gcm
>>> x, y = gcm.lonlat_to_xy(np.array([0.0, 5.0]), np.array([3.0, 4.0]))

tracmass.gcm.scrippsgog.xy_to_lonlat(x, y)[source]#

Convert fractional TRACMASS (x, y) positions to longitude/latitude.

Parameters:

x (array-like) – Fractional 0-based full-grid positions as stored in the trajectory Parquet file (columns "x" and "y").
y (array-like) – Fractional 0-based full-grid positions as stored in the trajectory Parquet file (columns "x" and "y").

Returns:

lon, lat (np.ndarray) – Longitude and latitude in degrees, same shape as x and y.

Return type:

tuple[ndarray, ndarray]

Examples

>>> df = pd.read_parquet("trm_gog.parquet", columns=["x", "y"])
>>> lon, lat = gcm.scrippsgog.xy_to_lonlat(df["x"], df["y"])