tracmass.integrate#

Core integration routines for TRACMASS-JAX.

Provides the inner lax.while_loop-based integration window and the two public integration functions that compute the full global flux grid:

steady_state() — single-window integration with a fixed flux field.
time_dependent() — multi-window integration with per-window field updates via update_flux().

For memory-efficient alternatives that avoid loading the full grid, see streaming.

tracmass.integrate.full_grid.steady_state(state, grid, fluxes, *, iter_sub=1, max_iter=100_000, timax_seconds=_UNSET, traj_max_time=_UNSET, scheme='tstep', no_wflux=_UNSET, n_particle_batches=_UNSET)[source]#

Single-window steady-state integration.

The flux direction (forward/backward) is determined by grid.reverse, as set when the grid was created via tracmass.gcm.new_grid() or tracmass.gcm.new_flux().

Parameters:

state (ParticleState) – Initial particle state.
grid (gcm.GridFields) – Static grid geometry from tracmass.gcm.new_grid() or tracmass.gcm.new_flux(). Pass dtm_end earlier than dtm_start to either for backward-in-time trajectories.
fluxes (gcm.FluxFields) – Time-varying flux data for the current time window.
iter_sub (int) – Number of sub-cycles per GCM step (ignored when scheme=’analytical’).
max_iter (int) – Maximum number of sub-cycles.
timax_seconds (float, optional) – Maximum trajectory lifetime in seconds, counted from each particle’s release (tt - t0). Pass this or traj_max_time, not both.
traj_max_time (str or None, optional) – The same lifetime as a pd.Timedelta string ("30D", "12h") or in calendar units ("1M", "6M", "1Q", "1Y", measured from the run’s start, fluxes.dtm). Particles past it stop with reason time. When neither is passed, from the traj_max_time setting (default: no limit). See tracmass.runsetup.traj_max_seconds().
scheme (str) – Advection scheme: "tstep" (default) or "analytical".
no_wflux (bool, optional) – When True, vertical (w) fluxes are ignored and particles are confined to their initial depth layer (ka never changes). When not passed, resolved from the move_vertically setting as no_wflux = not move_vertically (default move_vertically=True, i.e. no_wflux=False). See tracmass.config.move_vertically().
n_particle_batches (int, optional) – Split particles into this many equal-sized groups and run each group through lax.while_loop separately, capping the peak per-call GPU scratch allocation to roughly 1 / n_particle_batches. When not passed, taken from the n_particle_batches setting (default 1 — no batching); an explicit value wins. See tracmass.config.n_particle_batches() and time_dependent() for the full rationale.

Returns:

ParticleState with final positions.
For multi-window integration with field updates, use :func:`time_dependent.`

Return type:

ParticleState

tracmass.integrate.full_grid.time_dependent(state, grid, fluxes, fluxes_fn=None, *, iter_sub=1, max_iter=100_000, n_windows=None, max_sub_iter_per_window=10_000, timax_seconds=_UNSET, traj_max_time=_UNSET, recorder=_UNSET, monitor=_UNSET, scheme='tstep', no_wflux=_UNSET, pace_factor=5, n_particle_batches=_UNSET, rampup_windows=_UNSET, rampup_substeps=_UNSET, seed_freq=_UNSET)[source]#

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

Each window the flux buffer is advanced by one timestep via tracmass.gcm.update_flux(), which swaps time slots and loads the next flux snapshot from the fluxes_fn callable.

The GCM window duration is taken from grid.timestep; it does not need to be supplied separately.

Parameters:

state (ParticleState) – Initial particle state.
grid (gcm.GridFields) – Static grid geometry, as returned by tracmass.gcm.new_grid() or tracmass.gcm.new_flux(). grid.reverse controls the integration direction. grid.timestep sets the duration of each GCM window.
fluxes (gcm.FluxFields) – Initial flux data with both time slots populated.
fluxes_fn (callable or None) – Callable (e.g. gcm.theoretical.fluxes or gcm.glorys.fluxes) called as fluxes_fn(dtm, reverse=...) to load the next window. If None, inferred from fluxes.fluxes_fn when fluxes was created by new_flux(). Falls back to holding fields constant (steady-state flow repeated) if neither is set.
iter_sub (int) – Number of sub-cycles per GCM step (ignored when scheme=’analytical’).
max_iter (int) – Maximum number of sub-cycles.
n_windows (int or None, optional) – Number of GCM time windows to integrate over. Default None — derived from the flux buffer as fluxes.n_windows, i.e. enough windows to reach the dtm_end given to tracmass.gcm.new_flux(); 1 if no dtm_end was given.
max_sub_iter_per_window (int) – Safety cap on sub-cycles within each window. Default 10 000. If reached, stuck particles are deactivated and a RuntimeWarning is issued. Raise this value for very high-CFL flows (fast jets).
timax_seconds (float, optional) – Maximum trajectory lifetime in seconds, counted from each particle’s release (tt - t0). Pass this or traj_max_time, not both.
traj_max_time (str or None, optional) – The same lifetime as a pd.Timedelta string ("30D", "12h") or in calendar units ("1M", "6M", "1Q", "1Y", measured from the run’s start, fluxes.dtm). Particles past it stop with reason time. When neither is passed, from the traj_max_time setting (default: no limit). See tracmass.runsetup.traj_max_seconds().
recorder (ParticleRecorder or str or None, optional) – Where to stream active-particle snapshots (each window boundary) and deactivation events, all to one Parquet file (event_type column). Pass a ParticleRecorder to manage the file yourself, or a path string to have one built and closed for you. None disables recording. When omitted, a recorder is built from settings (output_name / auto_record / base_output_dir); if neither output_name nor auto_record is set, nothing is written. See tracmass.runsetup.
monitor (RunMonitor or str or None, optional) – Controls progress output. Pass a RunMonitor instance, or a backend name ("print", "log", "rich") to have one built for you. None suppresses all output. When omitted, the backend is taken from the runmonitor_backend setting (default "print"). A monitor / recorder built here is entered and exited automatically, so no with block is needed.
scheme (str) – "tstep" (default) or "analytical".
no_wflux (bool, optional) – When True, vertical (w) fluxes are ignored and particles are confined to their initial depth layer (ka never changes). When not passed, resolved from the move_vertically setting as no_wflux = not move_vertically (default move_vertically=True, i.e. no_wflux=False). See tracmass.config.move_vertically().
n_particle_batches (int, optional) –
Split particles into this many equal-sized groups and run each group through lax.while_loop separately. When omitted, taken from the n_particle_batches setting (default 1 — no batching); an explicit value wins. See tracmass.config.n_particle_batches().

Each lax.while_loop call pre-allocates a fixed block of GPU scratch memory proportional to the number of particles it processes. On GPUs with limited memory this single allocation can be the bottleneck even when total memory usage is not high. Setting n_particle_batches=4 reduces the peak per-call allocation to roughly 1/4, at the cost of running 4 while_loops per window instead of one. The flux arrays are kept as numpy and transferred to the GPU per batch (inside _integrate_window) and freed on return, so between batches the GPU holds nothing (no persistent flux, no carried-over state).

The only semantic change: the time-interpolation weight intrpg and the progressive-deactivation pace are computed per batch rather than globally. In practice all batches start at the same window boundary, so their mean_ts values are nearly identical.
rampup_windows (int, optional) – Save positions more often during the first rampup_windows windows: rampup_substeps times per window instead of once, i.e. every timestep / rampup_substeps. The extra saves are event_type="rampup" rows (window = the window they fall in, dtm = their exact time); the window-boundary ones stay "snapshot". The run metadata gets save_dt_seconds, the time between consecutive saves. Particles are stopped exactly on each save time and restarted, which also refreshes the time interpolation weight there: during ramp-up the run behaves like one with a finer iter_sub (more accurate, not bit-identical to a run without ramp-up). When omitted, from the rampup_windows setting (default 0: off).
rampup_substeps (int, optional) – Saves per window during ramp-up. When omitted, from the rampup_substeps setting (default 5). Not supported with scheme="analytical".
seed_freq (str or None, optional) – Release a copy of state again every seed_freq (a pandas frequency such as "30D" or "MS"), counted from the start of the run, at the start of the window it falls in (see tracmass.runsetup.seed_windows()). The copies start with tt = t0 = their release time and are appended after the existing particles, so they get new id values and are told apart by t0. The run metadata gets seed_freq and seed_windows. Each release grows the particle arrays, which recompiles the advection loop once. None: seed once. When omitted, from seed_freq in the [env.seed] settings table (default None).
pace_factor (int)

Returns:

ParticleState with final positions.

Return type:

ParticleState

Memory-efficient integration variants for TRACMASS-JAX.

Provides two integration functions that avoid loading the full global flux grid into device memory at once:

time_dependent_patched() — loads only the rectangular bounding box of the active particles (plus a configurable buffer) at each time window. The patch shape is fixed at initialisation so the JAX kernel is compiled once and reused.
time_dependent_streaming() — loads only the flux values at each particle’s current cell (a per-particle stencil) and advances particles with a Python for-loop instead of a JAX while_loop. Suitable for widely scattered particles or very long runs.

Both functions use _integrate_window() internally (patched variant) or replace it with an explicit Python loop (streaming variant).

tracmass.integrate.streaming.time_dependent_patched(state, grid, fluxes, fluxes_fn=None, *, iter_sub=1, max_iter=100_000, n_windows=None, max_sub_iter_per_window=10_000, timax_seconds=_UNSET, traj_max_time=_UNSET, recorder=_UNSET, monitor=_UNSET, no_wflux=_UNSET, patch_buffer=50, patch_granularity=100, pace_factor=5, n_particle_batches=_UNSET, rampup_windows=_UNSET, rampup_substeps=_UNSET, seed_freq=_UNSET)[source]#

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

Instead of computing the full global grid, this function computes the bounding box of the active particles at the start of each window, loads only the Zarr chunks that intersect that region plus patch_buffer cells of margin, and runs the advection kernel on the resulting sub-patch.

The patch shape is fixed at initialisation (initial bbox + patch_buffer on each side) so the JAX kernel is compiled exactly once and reused for all windows. Between windows the patch origin is recentred on the current particle centroid (zero-cost re-slicing of the lazy xarray DataArrays) and the patch content is re-computed from the appropriate Zarr chunks.

If particles drift far enough that the patch would need to grow beyond its initial size, those particles are allowed to reach the patch boundary and are treated as leaving the domain (deactivated). Use a generous patch_buffer to prevent this.

The GCM window duration is taken from grid.timestep; it does not need to be supplied separately.

Parameters:

state (ParticleState) – Initial particle state.
grid (gcm.GridFields) – Full-grid static geometry. grid.timestep sets the window duration.
fluxes (gcm.LazyFluxFields) – Lazy flux data for the first two time slots.
fluxes_fn (callable or None) – Called as fluxes_fn(dtm) to load the next window. If None, inferred from fluxes.fluxes_fn (set automatically by new_flux()).
iter_sub (int) – Sub-cycles per GCM step.
max_iter (int) – Maximum sub-cycles.
n_windows (int or None, optional) – Number of GCM time windows to integrate over. Default None — derived from the flux buffer as fluxes.n_windows, i.e. enough windows to reach the dtm_end given to tracmass.gcm.new_flux(); 1 if no dtm_end was given.
max_sub_iter_per_window (int) – Safety cap on sub-cycles per window. Default 10 000.
timax_seconds (float, optional) – Maximum trajectory lifetime in seconds, counted from each particle’s release (tt - t0). Pass this or traj_max_time, not both.
traj_max_time (str or None, optional) – The same lifetime as a pd.Timedelta string ("30D", "12h") or in calendar units ("1M", "6M", "1Q", "1Y", measured from the run’s start, fluxes.dtm). Particles past it stop with reason time. When neither is passed, from the traj_max_time setting (default: no limit). See tracmass.runsetup.traj_max_seconds().
recorder (ParticleRecorder or str or None, optional) – Where to stream active-particle snapshots (each window boundary) and deactivation events, all to one Parquet file (event_type column). Pass a ParticleRecorder to manage the file yourself, or a path string to have one built and closed for you. None disables recording. When omitted, a recorder is built from settings (output_name / auto_record / base_output_dir); if neither output_name nor auto_record is set, nothing is written. See tracmass.runsetup.
no_wflux (bool, optional) – Ignore vertical fluxes. When not passed, resolved from the move_vertically setting as no_wflux = not move_vertically (default move_vertically=True, i.e. no_wflux=False). See tracmass.config.move_vertically().
patch_buffer (int) – Extra cells added on every side of the particle bounding box.
patch_granularity (int, optional) – Patch dimensions are rounded up to the nearest multiple of this value. Limits the number of distinct patch shapes and therefore XLA recompilations as the particle cloud spreads. Default 100.
monitor (RunMonitor or str or None, optional) – Controls progress output. Pass a RunMonitor instance, or a backend name ("print", "log", "rich") to have one built for you. None suppresses all output. When omitted, the backend is taken from the runmonitor_backend setting (default "print"). A monitor / recorder built here is entered and exited automatically, so no with block is needed.
n_particle_batches (int, optional) – Accepted for API compatibility with the full-grid integrators (steady_state() / time_dependent()) — resolved from the n_particle_batches setting when omitted — but has no effect here. Default 1.
rampup_windows (int, optional) – Extra event_type="rampup" saves in the first windows, as in time_dependent().
rampup_substeps (int, optional) – Extra event_type="rampup" saves in the first windows, as in time_dependent().
seed_freq (str or None, optional) – Release a copy of state again every seed_freq (pandas frequency, e.g. "30D" or "MS"), as in time_dependent(). When omitted, from seed_freq in the [env.seed] settings table.
pace_factor (int)

Returns:

ParticleState with final positions in **global** grid coordinates.

Return type:

ParticleState

tracmass.integrate.streaming.time_dependent_streaming(state, grid, fluxes, fluxes_fn=None, *, iter_sub=1, max_iter=100_000, n_windows=None, max_sub_iter_per_window=10_000, timax_seconds=_UNSET, traj_max_time=_UNSET, recorder=_UNSET, monitor=_UNSET, no_wflux=_UNSET, pace_factor=5, n_particle_batches=_UNSET, rampup_windows=_UNSET, rampup_substeps=_UNSET, seed_freq=_UNSET)[source]#

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

At each sub-cycle, only the flux values needed for each particle’s current cell are loaded via vectorised xarray point indexing — no rectangular patch is pre-allocated. This makes the approach viable for widely scattered particles and very long runs where a bounding-box patch would cover most of the global grid.

The JAX while_loop used by time_dependent() and time_dependent_patched() is replaced by a Python for-loop, with one Dask compute and one JAX JIT call per sub-cycle. The total number of Python iterations is n_windows × (subcycles per window), typically on the order of n_windows × 10–20 for monthly GLORYS windows.

The GCM window duration is taken from grid.timestep; it does not need to be supplied separately.

Parameters:

state (ParticleState) – Initial particle state.
grid (gcm.GridFields) – Static grid geometry. grid.timestep sets the window duration.
fluxes (gcm.LazyFluxFields) – Lazy flux data for the first two time slots.
fluxes_fn (callable or None) – Called as fluxes_fn(dtm) to load the next window. If None, inferred from fluxes.fluxes_fn (set automatically by new_flux()).
iter_sub (int) – Sub-cycles per GCM step; sets minimum sub-cycle duration grid.timestep / iter_sub.
max_iter (int) – Maximum number of sub-cycles over the entire run.
n_windows (int or None, optional) – Number of GCM time windows to integrate over. Default None — derived from the flux buffer as fluxes.n_windows, i.e. enough windows to reach the dtm_end given to tracmass.gcm.new_flux(); 1 if no dtm_end was given.
max_sub_iter_per_window (int) – Safety cap on sub-cycles within each window. Default 10 000.
timax_seconds (float, optional) – Maximum trajectory lifetime in seconds, counted from each particle’s release (tt - t0). Pass this or traj_max_time, not both.
traj_max_time (str or None, optional) – The same lifetime as a pd.Timedelta string ("30D", "12h") or in calendar units ("1M", "6M", "1Q", "1Y", measured from the run’s start, fluxes.dtm). Particles past it stop with reason time. When neither is passed, from the traj_max_time setting (default: no limit). See tracmass.runsetup.traj_max_seconds().
recorder (ParticleRecorder or str or None, optional) – Where to stream active-particle snapshots (each window boundary) and deactivation events, all to one Parquet file (event_type column). Pass a ParticleRecorder to manage the file yourself, or a path string to have one built and closed for you. None disables recording. When omitted, a recorder is built from settings (output_name / auto_record / base_output_dir); if neither output_name nor auto_record is set, nothing is written. See tracmass.runsetup.
monitor (RunMonitor or str or None, optional) – Controls progress output. Pass a RunMonitor instance, or a backend name ("print", "log", "rich") to have one built for you. None suppresses all output. When omitted, the backend is taken from the runmonitor_backend setting (default "print"). A monitor / recorder built here is entered and exited automatically, so no with block is needed.
no_wflux (bool, optional) – When True, vertical (w) fluxes are ignored and particles are confined to their initial depth layer. When not passed, resolved from the move_vertically setting as no_wflux = not move_vertically (default move_vertically=True, i.e. no_wflux=False). See tracmass.config.move_vertically().
n_particle_batches (int, optional) – Accepted for API compatibility with the full-grid integrators (steady_state() / time_dependent()) — resolved from the n_particle_batches setting when omitted — but has no effect here. The streaming variant already processes all particles in a single Python loop without a JAX while_loop, so there is no per-call GPU scratch allocation to reduce. Default 1.
rampup_windows (int, optional) – Not supported here: a ramp-up (from the arguments or the rampup_windows setting) raises NotImplementedError; use time_dependent() or time_dependent_patched().
rampup_substeps (int, optional) – Not supported here: a ramp-up (from the arguments or the rampup_windows setting) raises NotImplementedError; use time_dependent() or time_dependent_patched().
seed_freq (str or None, optional) – Release a copy of state again every seed_freq (pandas frequency, e.g. "30D" or "MS"), as in time_dependent(). When omitted, from seed_freq in the [env.seed] settings table.
pace_factor (int)

Returns:

ParticleState – Final particle state in global grid coordinates.

Return type:

ParticleState