API reference#

Generated from the NumPy-style docstrings in src/tracmass. Every public module, class and function with a docstring is included; private helpers (names starting with _) are left out.

TRACMASS-JAX: Lagrangian particle tracking on JAX-accelerated GCM grids.

This package provides the core data structures, integration routines, and GCM interfaces for running Lagrangian particle trajectories on ocean model velocity fields, using JAX for hardware-accelerated computation.

The grid is split into two objects so that flux data can be loaded lazily:

  • GridFields — static geometry (kmt, dxdy) and metadata. Closed over by compiled kernels as XLA constants.

  • FluxFields — time-varying flux data for a two-slot window. Passed as a dynamic argument; load it just before calling the kernel.

Typical usage#

>>> import tracmass
>>> grid, fluxes = tracmass.new_flux(tracmass.gcm.theoretical.fluxes)
>>> state = tracmass.gcm.theoretical.initial_state()
>>> result = tracmass.steady_state(state, grid, fluxes)

new_flux builds the grid and the flux buffer together. To reuse one grid across several flux buffers, build it once with new_grid() and pass it back in:

>>> grid = tracmass.new_grid(tracmass.gcm.theoretical.fluxes)
>>> fluxes = tracmass.new_flux(grid, tracmass.gcm.theoretical.fluxes)

Top-level functions#

run

Build the grid and flux buffer, seed, and integrate in one call.

new_grid

Build the static grid geometry for a run.

new_flux

Create the lazy two-slot flux buffer for a run.

update_flux

Advance the flux buffer by one timestep using a rolling two-slot window.

swap_fields

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

steady_state

Single-window steady-state integration.

time_dependent

Multi-window integration using the ned (new/update grid) system.

time_dependent_patched

Multi-window integration using local flux patches around the particles.

time_dependent_streaming

Multi-window integration using per-particle flux stencil extraction.

load

Load a trajectory file written by ParticleRecorder.

write_trajectories_parquet

Write a single particle state snapshot to a Parquet file.

set_verbose

Turn verbose debug output on or off for the session.

free_device_memory

Release GPU/device memory held by tracmass between runs.

Classes#

ParticleState

State vector for all particles — each field is shape (n,).

GridFields

Static grid geometry and metadata for a TRACMASS run.

LazyFluxFields

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

FluxFields

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

ParticleRecorder

Streams trajectory snapshots and deactivation events to a Parquet file.

RunMonitor

Monitor for time_dependent() runs.

Reason codes#

REASON_ACTIVE (0), REASON_TIME (1), REASON_LAND (2), REASON_BOUNDS_H (3), REASON_PROGRESS (4), REASON_SLOW (5), REASON_ITERLIMIT (6), REASON_DXYZ (7), REASON_BOUNDS_V (8). See Stopping trajectories.

Modules#