tracmass.gcm.flux_grid#

Grid field containers and utilities for TRACMASS JAX.

This module defines the core data structures and helper functions used to manage GCM grid geometry and velocity flux fields across two-slot time windows required by the TRACMASS Lagrangian particle tracker.

The grid is split into three objects:

GridFields — static geometry (kmt, dxdy) and metadata (dimensions, periodicity). These fields never change during a run and are closed over by the compiled advection kernels as XLA constants.
LazyFluxFields — two time-slots of flux data held as raw xarray DataArrays (possibly Dask-backed), with no I/O triggered yet. Also carries dtm, the timestamp of the past (slot 0) time level. Call compute() to load data into device memory just before calling the advection kernel. update_flux() advances the window by a pure Python dict reference swap — no Dask compute, no JAX, no GC required until the next compute().
FluxFields — the computed form: JAX arrays backed by device memory, registered as a pytree so the compiled kernel can receive it as a dynamic argument.

class tracmass.gcm.flux_grid.GridFields(*, kmt, dxdy, imt, jmt, km, timestep, reverse, iperio, jperio, i_offset=0, j_offset=0)[source]#

Bases: object

Static grid geometry and metadata for a TRACMASS run.

These fields are fixed for the lifetime of a run. The advection kernels close over a GridFields instance so that kmt and dxdy become XLA constants in the compiled graph.

Parameters:

kmt (Array)
dxdy (Array)
imt (int)
jmt (int)
km (int)
timestep (Timedelta)
reverse (bool)
iperio (int)
jperio (int)
i_offset (int)
j_offset (int)

kmt#

Bottom-level index (integer), shape (imt, jmt). A value of 0 indicates a land cell.

Type:: jnp.ndarray

dxdy#

Horizontal cell area, shape (imt, jmt).

Type:: jnp.ndarray

imt#

Number of grid cells in the zonal (x) direction.

Type:: int

jmt#

Number of grid cells in the meridional (y) direction.

Type:: int

km#

Number of vertical levels.

Type:: int

timestep#

Duration between consecutive time levels.

Type:: pd.Timedelta

reverse#

If True, fluxes have been negated for backward-in-time tracking.

Type:: bool

iperio#

East-west periodicity flag (1 = cyclic, 0 = closed).

Type:: int

jperio#

North-fold periodicity flag (1 = north-fold, 0 = closed).

Type:: int

i_offset, j_offset

Position of this grid’s cell (0, 0) in the adapter’s full (un-subset) grid — the subset i1 / j1, else 0. Particle positions are relative to this grid; the trajectory recorder adds the offsets so saved x / y / ia / ja are absolute.

Type:: int

tracmass.gcm.flux_grid.edge_mask(grid)[source]#

Cells on the edge of the domain, where particles leave it.

The outer ring of cells along every axis that does not wrap around: the first and last column when grid.iperio == 0 (e.g. a subset in longitude) and the first and last row when grid.jperio == 0. A particle that enters one of these cells stops with REASON_BOUNDS_H. The edge cells cannot be integrated correctly: the flux through their outer face is not part of the grid (the kernels treat the west / south one as a closed wall, and a particle crossing the east / north one would read fluxes outside the arrays).

Parameters:: grid (GridFields) – Grid geometry; only imt, jmt, iperio and jperio are used.
Returns:: np.ndarray – Boolean array, shape (imt, jmt).
Return type:: ndarray

Examples

>>> edge = edge_mask(grid)            # subset grid: iperio == 0
>>> edge[0].all() and edge[-1].all()
True

class tracmass.gcm.flux_grid.FluxFields(*, dzt, zstot, uflux, vflux, wflux=None)[source]#

Bases: object

Two-slot time window of velocity flux data for the advection kernels.

Each array field has a trailing dimension of size 2, where slot 0 holds the ‘past’ time level and slot 1 holds the ‘present’ time level.

FluxFields is a JAX pytree and is passed as a dynamic argument to the compiled step functions. Obtain one from a LazyFluxFields by calling compute().

Parameters:

dzt (Array)
zstot (Array)
uflux (Array)
vflux (Array)
wflux (Array | None)

dzt#

Layer thickness at tracer points, shape (imt, jmt, km, 2).

Type:: jnp.ndarray

zstot#

Total sea-surface height anomaly, shape (imt, jmt, 2).

Type:: jnp.ndarray

uflux#

Zonal volume flux, shape (imt, jmt, km, 2).

Type:: jnp.ndarray

vflux#

Meridional volume flux, shape (imt, jmt+1, km, 2).

Type:: jnp.ndarray

wflux#

Vertical volume flux in the kernel convention: wflux[..., k, :] is the flux through the bottom face of level k (level 0 at the surface), positive downward. Adapters supply the model convention (top face, positive upward); see _w_to_kernel(). None disables vertical particle motion.

Type:: jnp.ndarray or None

class tracmass.gcm.flux_grid.LocalFluxFields(*, uflux, vflux, wflux, dzt, zstot, kmt, dxdy, edge=None)[source]#

Bases: object

Flux data for a spatial sub-patch, with geometry included dynamically.

Unlike FluxFields, kmt and dxdy are carried as ordinary JAX arrays rather than being closed over as XLA constants. This allows the advection kernel to be compiled once for a fixed patch shape and then reused across windows as the patch content changes — the kernel receives new values via the normal dynamic-argument path.

Parameters:

uflux (Array)
vflux (Array)
wflux (Array | None)
dzt (Array)
zstot (Array)
kmt (Array)
dxdy (Array)
edge (Array | None)

uflux#

Zonal volume flux, shape (imt_p, jmt_p, km, 2).

Type:: jnp.ndarray

vflux#

Meridional volume flux, shape (imt_p, jmt_p+1, km, 2).

Type:: jnp.ndarray

wflux#

Vertical volume flux, shape (imt_p, jmt_p, km+1, 2).

Type:: jnp.ndarray or None

dzt#

Layer thickness, shape (imt_p, jmt_p, km, 2).

Type:: jnp.ndarray

zstot#

Sea-surface height anomaly, shape (imt_p, jmt_p, 2).

Type:: jnp.ndarray

kmt#

Bottom-level index, shape (imt_p, jmt_p). Dynamic — values differ per patch but shape is fixed for the lifetime of the run.

Type:: jnp.ndarray

dxdy#

Horizontal cell area, shape (imt_p, jmt_p). Dynamic.

Type:: jnp.ndarray

edge#

edge_mask() of the full grid sliced to the patch, shape (imt_p, jmt_p): the patch cells on the domain edge, where particles stop as leaving the domain. None: no domain edge in the patch. Dynamic.

Type:: jnp.ndarray or None

class tracmass.gcm.flux_grid.FluxStencil(*, u_east, u_west, v_north, v_south, w_top, w_bot, dzt, zstot, kmt, dxdy)[source]#

Bases: object

Per-particle flux stencil extracted at each particle’s current cell.

Each array carries values for all N active particles, with a trailing slot dimension of size 2 (slot 0 = past, slot 1 = present).

Obtained via extract_stencil() using vectorised xarray point indexing, which tells Dask to load only the Zarr chunks that contain the requested grid cells. This enables memory-efficient integration of widely scattered particles without loading a rectangular spatial patch.

Parameters:

u_east (Array)
u_west (Array)
v_north (Array)
v_south (Array)
w_top (Array | None)
w_bot (Array | None)
dzt (Array)
zstot (Array)
kmt (Array)
dxdy (Array)

u_east#

Zonal flux at the east wall of each particle’s cell, shape (N, 2).

Type:: jnp.ndarray

u_west#

Zonal flux at the west wall, shape (N, 2).

Type:: jnp.ndarray

v_north#

Meridional flux at the north wall, shape (N, 2).

Type:: jnp.ndarray

v_south#

Meridional flux at the south wall, shape (N, 2).

Type:: jnp.ndarray

w_top#

Vertical flux at the top wall, shape (N, 2). None when vertical motion is disabled (no_wflux=True).

Type:: jnp.ndarray or None

w_bot#

Vertical flux at the bottom wall, shape (N, 2) or None.

Type:: jnp.ndarray or None

dzt#

Layer thickness at the particle’s cell, shape (N, 2).

Type:: jnp.ndarray

zstot#

Total column depth at the particle’s position, shape (N, 2).

Type:: jnp.ndarray

kmt#

Bottom level index at the particle’s cell, shape (N,).

Type:: jnp.ndarray

dxdy#

Horizontal cell area at the particle’s position, shape (N,).

Type:: jnp.ndarray

tracmass.gcm.flux_grid.extract_stencil(lazy, grid, state, no_wflux=False, mask=None)[source]#

Extract per-particle flux stencils using vectorised xarray point indexing.

For each particle at grid position (ia, ja, ka) (1-based), reads the six flux face values required for one advection sub-cycle (east/west for u, north/south for v, top/bottom for w) plus cell geometry (dzt, zstot, kmt, dxdy). All xarray DataArrays are fetched in a single dask.compute() call so Dask loads only the Zarr chunks that contain the requested points.

Parameters:

lazy (LazyFluxFields) – Lazy flux data for the current time window.
grid (GridFields) – Static grid geometry (provides kmt, dxdy, and iperio).
state (ParticleState) – Current particle state. Only ia, ja, ka are used.
no_wflux (bool, optional) – When True, w_top and w_bot are set to None.
mask (np.ndarray or None, optional) – Boolean array of shape (N,) selecting which particles to read Zarr data for. Positions where mask is False receive zeros in the returned FluxStencil (those values are never used by the kernel’s jnp.where(state.active, ...) guards). When None (default), data is read only for active particles (state.active), skipping unnecessary I/O for deactivated ones.

Returns:

FluxStencil – Per-particle flux stencils, shape (N, 2) / (N,). Only positions covered by mask (or state.active when mask is None) carry valid data; all other positions are zero.

Return type:

FluxStencil

class tracmass.gcm.flux_grid.LazyFluxFields(_past, _present, dtm=None, fluxes_fn=None, zarr_chunk_xy=(1, 1), timestep=None, reverse=False, dtm_end=None, _computed=<factory>)[source]#

Bases: object

Two time-slots of flux data held as raw xarray DataArrays.

The DataArrays may be backed by Dask (for GLORYS from Zarr) or by in-memory arrays (for the theoretical adapter). In either case, no numpy or device memory is allocated until compute() is called.

This object is not a JAX pytree. Pass the result of compute() to the advection kernel instead.

The two time slots are stored as separate DataArrays without a trailing time dimension. The shape of each DataArray matches the spatial shape of the variable (e.g. (imt, jmt, km) for uflux). update_flux() advances the window by reassigning Python dict references — no Dask compute, no JAX, no GC.

Parameters:

_past (dict)
_present (dict)
dtm (Timestamp | None)
fluxes_fn (Callable | None)
zarr_chunk_xy (tuple[int, int])
timestep (Timedelta | None)
reverse (bool)
dtm_end (Timestamp | None)
_computed (dict)

_past#

Mapping {name: xr.DataArray} for the ‘past’ time slot.

Type:: dict

_present#

Mapping {name: xr.DataArray} for the ‘present’ time slot.

Type:: dict

dtm#

Timestamp of the ‘past’ (slot 0) time level. Set by new_flux() and update_flux(); None when constructed via lower-level helpers like swap_fields().

Type:: pd.Timestamp or None

fluxes_fn#

The GCM adapter callable used to create this object, e.g. gcm.glorys.fluxes. Stored by new_flux() and propagated by update_flux() so that time_dependent() can infer it automatically without the caller re-supplying it.

Type:: callable or None

timestep#

GCM output interval (always positive), copied from grid.timestep by new_flux(). update_flux() steps dtm by this amount.

Type:: pd.Timedelta or None

reverse#

True for a backward run (copied from grid.reverse): update_flux() then walks dtm backward and negates the velocity fluxes it loads.

Type:: bool

dtm_end#

The dtm_end passed to new_flux(), or None. Used by n_windows to size a run.

Type:: pd.Timestamp or None

Notes

compute() keeps the present slot’s arrays (already read and cleaned) in _computed, keyed by the slot’s DataArray. After update_flux() that DataArray is the new past slot, so the next compute() reuses the arrays instead of reading the same snapshot again: every snapshot is read once, not twice. _computed is shared by the copies update_flux() makes and holds one slot per variable.

property n_windows: int | None#

Number of GCM windows left between dtm and dtm_end.

ceil(|dtm_end - dtm| / timestep), so a dtm_end that falls inside a window still has that window integrated. Decreases by one with every update_flux() call. None when dtm_end (or dtm / timestep) is not set.

Returns:: int or None

compute_patch(i0, i1, j0, j1, grid, include_wflux=True)[source]#

Load only the sub-patch [i0:i1, j0:j1] into device memory.

Slices the lazy xarray DataArrays before calling .values, so Dask only computes (and reads from Zarr) the chunks that intersect the requested region. kmt and dxdy are sliced from grid and included in the returned LocalFluxFields as dynamic JAX arrays, so the advection kernel can be compiled once for a fixed patch shape and reused without recompilation as patch content changes.

Parameters:

i0 (int) – Zonal slice bounds (0-based, exclusive end).
i1 (int) – Zonal slice bounds (0-based, exclusive end).
j0 (int) – Meridional slice bounds (0-based, exclusive end).
j1 (int) – Meridional slice bounds (0-based, exclusive end).
grid (GridFields) – Full-grid static geometry; kmt and dxdy are sliced from here to match the patch.
include_wflux (bool, optional) – False skips reading wflux (returned as None), for runs without vertical motion. Default True.

Returns:

LocalFluxFields

Return type:

LocalFluxFields

preload(include_wflux=True)[source]#

Return a copy with all Dask arrays preloaded to numpy.

Issues a single dask.compute() call for every variable in _past and _present, replacing Dask-backed DataArrays with in-memory numpy-backed DataArrays. The returned object is otherwise identical and can be used anywhere a LazyFluxFields is expected.

Calling this once per GCM window in time_dependent_streaming() changes the I/O pattern from O(n_sub) small Dask computes (one per sub-cycle that has cell crossings) to O(1) per window. Subsequent extract_stencil() calls on the returned object use pure numpy fancy indexing — no further Zarr I/O regardless of how many sub-cycles or cell crossings occur.

Parameters:: include_wflux (bool, optional) – False leaves wflux out of the copy instead of reading it, for runs without vertical motion. Default True.
Returns:: LazyFluxFields – A copy with numpy-backed DataArrays and the same metadata (dtm, fluxes_fn, timestep, reverse, …).
Return type:: LazyFluxFields

compute(include_wflux=True)[source]#

Trigger Dask computation and transfer flux data to JAX device memory.

Variables are loaded one at a time (past then present) and stacked along a new trailing axis, so the peak extra allocation per variable is one numpy array of size t0, one of size t1, and the resulting JAX array of size t0 + t1.

NaN values in flux arrays (uflux, vflux, wflux) are replaced with 0.0 before stacking. Land cells in GLORYS/NEMO are masked as NaN in the Zarr files; the physical no-flow boundary condition is 0.0, so NaN → 0 is the correct substitution. Without this replacement, NaN fluxes propagate through _calc_pos_1d to produce a NaN particle position, which after astype(int32) conversion gives a garbage cell index that may point to a kmt=0 (land) cell, causing spurious REASON_LAND deactivations.

Parameters:: include_wflux (bool, optional) – False skips reading wflux (returned as None), for runs without vertical motion. Default True.
Returns:: FluxFields – Kernel-ready flux data with trailing slot dimension of size 2.
Return type:: FluxFields

tracmass.gcm.flux_grid.subset_offset_attrs(subset)[source]#

Dataset attrs recording where an adapter’s x/y subset starts.

Subsetting adapters put these on the dataset fluxes() returns; _static_from_ds() turns them into GridFields.i_offset / j_offset so recorded positions can be written in full-grid indices.

Parameters:: subset (dict) – The adapter’s settings["subset"] (i1, j1, istep, …).
Returns:: dict[str, int] – {"_i_offset": i1 or 0, "_j_offset": j1 or 0}.
Warns:: UserWarning – If istep / jstep is set: a strided subset’s cells are not full-grid cells, so the offset alone does not make saved positions absolute.
Return type:: dict[str, int]

tracmass.gcm.flux_grid.swap_fields(lazy_fluxes)[source]#

Swap time levels: present (slot 1) becomes past (slot 0).

Mirrors mod_swap.F90::swap_time(). After calling this function, the caller should provide new present-slot data (e.g. by calling update_flux()).

This is a pure Python reference swap — no Dask compute, no JAX.

Parameters:: lazy_fluxes (LazyFluxFields) – Lazy flux data to swap.
Returns:: LazyFluxFields – New LazyFluxFields where the old present data is now the past slot. The present slot is left pointing at the old present data (stale) and should be overwritten by the caller.
Return type:: LazyFluxFields

tracmass.gcm.flux_grid.fluxes_to_grid(ds)[source]#

Convert a two-slot concatenated xarray Dataset to (GridFields, LazyFluxFields).

Static variables (kmt, dxdy, dimensions, timestamps) are computed eagerly into a GridFields. Flux variables (uflux, vflux, wflux, dzt, zstot) are kept as xarray DataArrays inside a LazyFluxFields — no I/O is triggered for them until compute() is called.

The DataArrays stored in the LazyFluxFields are expected to already have a trailing time slot dimension of size 2 (as produced by xr.concat). new_flux() builds the same two-slot structure directly from two separate fluxes() calls.

Parameters:: ds (xr.Dataset) – Two-slot concatenated dataset.
Returns:: tuple[GridFields, LazyFluxFields]
Return type:: tuple[GridFields, LazyFluxFields]

tracmass.gcm.flux_grid.new_grid(fluxes, dtm_start=None, dtm_end=None)[source]#

Build the static grid geometry for a run.

Calls fluxes once (at dtm_start) purely to read the time-invariant geometry — kmt, dxdy, imt / jmt / km, iperio / jperio — and the GCM output interval (timestep). No flux data is loaded. The returned GridFields is immutable for the lifetime of the run and can be reused across many new_flux() buffers on the same grid (successive experiments, restart windows).

Integration direction is inferred from dtm_start / dtm_end exactly as in new_flux(): dtm_end < dtm_start marks a backward run (grid.reverse=True). grid.timestep is always positive.

Parameters:

fluxes (callable or str) – fluxes(dtm) -> xr.Dataset, or the name of a tracmass.gcm sub-module (e.g. "glorys"). When the adapter can be identified, the [env.gcm] subset table from the active tracmass.toml environment is applied first via tracmass.config.apply_gcm_subset().
dtm_start (pd.Timestamp or str or None, optional) – Any valid time for the adapter; only the geometry is read. Default None — fluxes is called with no arguments at all, so its own default applies (for gcm.glorys/gcm.ecco2 this falls through to tracmass.seagrid.glorys.open_dataset() / tracmass.seagrid.ecco2.open_dataset()’s own dtm default).
dtm_end (pd.Timestamp or str or None, optional) – When earlier than the resolved dtm_start, marks the grid backward (reverse=True). Default None (forward).

Returns:

GridFields – Static geometry with timestep (positive) and reverse set.

Return type:

GridFields