"""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:
- :func:`steady_state` — single-window integration with a fixed flux field.
- :func:`time_dependent` — multi-window integration with per-window field
updates via :func:`~tracmass.gcm.update_flux`.
For memory-efficient alternatives that avoid loading the full grid, see
:mod:`~tracmass.integrate.streaming`.
"""
from __future__ import annotations
import functools
import gc
import logging
import time
import warnings
from typing import Callable
import jax
import jax.numpy as jnp
import numpy as np
from jax import lax
from tracmass import advection_scheme, gcm
from tracmass.gcm.flux_grid import FluxFields, LazyFluxFields
from tracmass.monitoring import RunMonitor
from tracmass.verbose import timed
from tracmass.particle import (
_REASON_NAMES,
REASON_ITERLIMIT,
REASON_PROGRESS,
REASON_SLOW,
REASON_TIME,
ParticleRecorder,
ParticleState,
)
from tracmass.runsetup import (
_UNSET,
manage_run_io,
rampup_metadata,
resolve_n_particle_batches,
resolve_n_windows,
resolve_no_wflux,
resolve_rampup,
resolve_seed_freq,
resolve_timax,
reseed,
seed_metadata,
seed_windows,
window_ends,
)
log = logging.getLogger(__name__)
def _ensure_computed(fluxes, no_wflux: bool = False) -> FluxFields:
"""Return a concrete ``FluxFields``, triggering Dask I/O if needed.
With *no_wflux*, ``wflux`` is not read and is ``None`` in the result.
"""
if isinstance(fluxes, LazyFluxFields):
return fluxes.compute(include_wflux=not no_wflux)
if no_wflux and fluxes.wflux is not None:
return fluxes._replace(wflux=None)
return fluxes
def _integrate_batched(
state: ParticleState,
step_fn,
fluxes,
*,
n_batches: int,
tseas: float,
iter_sub: int,
max_sub_iter: int,
timax_seconds: float,
ts_window_end: float,
pace_factor: int,
ts_window_start: float | None = None,
) -> tuple[ParticleState, int]:
"""Run ``_integrate_window`` in *n_batches* particle slices.
Splitting the population reduces the peak GPU allocation inside each
``lax.while_loop`` call by ~1/n_batches.
``fluxes`` is passed as numpy (CPU) arrays. ``_integrate_window``
transfers them to device at entry and they are freed when that function
returns, so **between batches the GPU holds nothing** (no persistent flux,
no carried-over state). Peak GPU per batch is:
``flux (H2D) + one batch state (H2D) + XLA scratch``.
Parameters
----------
state : ParticleState
Full particle state with **numpy** (CPU) arrays, shape ``(n,)`` per
field. The caller must D2H-convert and free the GPU copy before
calling so the GPU starts the batch loop empty.
step_fn : callable
Compiled advection kernel.
fluxes : FluxFields
Flux data as **numpy** arrays (CPU). Transferred to device per batch
inside ``_integrate_window`` and freed on return.
n_batches : int
Number of equal-sized particle slices.
tseas, iter_sub, max_sub_iter, timax_seconds, ts_window_end, pace_factor
Forwarded unchanged to ``_integrate_window``.
ts_window_start : float or None, optional
Forwarded unchanged to ``_integrate_window``.
Returns
-------
tuple[ParticleState, int]
Reassembled full-population state (JAX device arrays) and the maximum
sub-cycle count seen across all batches.
"""
n = len(state.active)
batch_size = (n + n_batches - 1) // n_batches
batch_results: list[ParticleState] = []
max_n_sub = 0
for b in range(n_batches):
lo = b * batch_size
hi = min(lo + batch_size, n)
if lo >= n:
break
# Slice on CPU (zero-copy numpy view), put to device, run, return to
# CPU. The GPU holds at most: flux + batch + XLA scratch at any time.
state_b_np = jax.tree_util.tree_map(lambda a, lo=lo, hi=hi: a[lo:hi], state)
state_b = jax.device_put(state_b_np)
# fluxes is numpy; _integrate_window's internal device_put transfers
# it to GPU and the local reference is dropped on return, freeing the
# GPU flux before the next batch starts.
state_b, n_sub_b = _integrate_window(
state_b,
step_fn,
fluxes,
tseas=tseas,
iter_sub=iter_sub,
max_sub_iter=max_sub_iter,
timax_seconds=timax_seconds,
ts_window_end=ts_window_end,
pace_factor=pace_factor,
ts_window_start=ts_window_start,
)
max_n_sub = max(max_n_sub, n_sub_b)
# D2H result; free GPU batch immediately so the next batch starts
# with a clean device.
batch_results.append(jax.tree_util.tree_map(np.asarray, state_b))
del state_b
gc.collect()
# Flush all pending XLA async operations (including deferred cudaFree
# callbacks) so GPU memory is fully reclaimed before the next batch
# allocates its state + flux + XLA scratch.
jax.effects_barrier()
state_np = jax.tree_util.tree_map(
lambda *arrs: np.concatenate(arrs, axis=0), *batch_results
)
return state_np, max_n_sub # numpy — caller keeps state on CPU between windows
# ===========================================================================
# Inner integration loop (within one GCM time window)
# ===========================================================================
# Fraction of a sub-step within which a particle counts as already on a
# sub-step boundary (absorbs float32 rounding of tt / ts; 86 s for daily
# fields with iter_sub=1).
_SUBSTEP_TOL = 1e-3
def _substep_cap(ts, tseas: float, iter_sub: int, ts_end=None):
"""Time left until each particle's next sub-step boundary.
Sub-step boundaries are at ``ts = n / iter_sub``; window boundaries are
the integer ones. Capping every step at this time (as Fortran TRACMASS
does) makes particles stop *on* the window boundary instead of
finishing a cell crossing that ends past it. With *ts_end* the cap is
also never past ``ts_end``, so particles stop exactly on a ramp-up save
time inside a window (a no-op for whole-window ends, which are sub-step
boundaries already).
Parameters
----------
ts : jnp.ndarray
Normalised particle time (``tt / tseas``), shape ``(N,)``.
tseas : float
GCM window length in seconds.
iter_sub : int
Sub-steps per window.
ts_end : jnp.ndarray or None, optional
Integration end (normalised time); caps ``ts_next``.
Returns
-------
ts_next : jnp.ndarray
``ts`` of the next boundary, shape ``(N,)``.
cap_s : jnp.ndarray
Seconds until it — the step cap passed to the kernel as ``dsmin``.
"""
n = jnp.floor(ts * iter_sub + _SUBSTEP_TOL) + 1.0
ts_next = n / iter_sub
if ts_end is not None:
ts_next = jnp.minimum(ts_next, ts_end)
return ts_next, (ts_next - ts) * tseas
def _snap_to_substep(st, tt_before, ts_next, cap_s, tseas: float, iter_sub: int):
"""Set ``ts = tt / tseas``, snapping time-capped steps onto the boundary.
A particle whose step used (within tolerance) exactly its *cap_s*
reached the sub-step boundary; its ``tt`` / ``ts`` are set to it
exactly, so float32 rounding cannot leave it a hair short of the window
end. Steps longer than the cap (the analytical scheme ignores it and
caps at the window end itself) are left alone.
"""
tol = _SUBSTEP_TOL * tseas / iter_sub
dt = st.tt - tt_before
hit = (jnp.abs(dt - cap_s) <= tol) & (dt > 0)
return st._replace(
tt=jnp.where(hit, ts_next * tseas, st.tt),
ts=jnp.where(hit, ts_next, st.tt / tseas),
)
@functools.partial(
jax.jit,
static_argnames=(
"step_fn",
"tseas",
"iter_sub",
"max_sub_iter",
"timax_seconds",
"pace_factor",
),
)
def _run_window(
state: ParticleState,
fluxes,
ts_end: jnp.ndarray,
ts_start: jnp.ndarray,
*,
step_fn,
tseas: float,
iter_sub: int,
max_sub_iter: int,
timax_seconds: float,
pace_factor: int,
) -> tuple[ParticleState, jnp.ndarray, jnp.ndarray]:
"""JIT-compiled while-loop for a single GCM window.
``fluxes`` is a **dynamic** JIT argument so JAX traces it with abstract
values (shape + dtype only). This guarantees a single XLA compilation
per ``(step_fn, tseas, …)`` combination regardless of the concrete flux
data, preventing per-window CUBIN accumulation in GPU VRAM. So are
``ts_end`` / ``ts_start``: a ramp-up window integrated in pieces
(``ts_start .. ts_end`` shorter than one window) reuses the program.
Returns
-------
tuple[ParticleState, jnp.ndarray, jnp.ndarray]
(state, n_sub, n_stuck) — all JAX arrays; caller D2H syncs as needed.
"""
_N_WARMUP = 2
_pace_denom = max((max_sub_iter - _N_WARMUP) // pace_factor, 1)
def cond_fn(carry):
st, i, ts_end_, _ = carry
any_active = jnp.any(st.active)
all_past_window = jnp.all(~st.active | (st.ts >= ts_end_))
return any_active & ~all_past_window & (i < max_sub_iter)
def body_fn(carry):
st, i, ts_end_, ts_start_ = carry
exceeded = (st.tt - st.t0) > timax_seconds
past_window = st.ts >= ts_end_
st = st._replace(
reason=jnp.where(st.active & exceeded, REASON_TIME, st.reason),
active=st.active & ~exceeded & ~past_window,
)
n_active = jnp.sum(st.active)
mean_ts = jnp.where(
n_active > 0,
jnp.sum(jnp.where(st.active, st.ts, 0.0)) / jnp.maximum(n_active, 1),
0.0,
)
intrpg = jnp.mod(mean_ts, 1.0)
tt_before = st.tt
ts_next, cap_s = _substep_cap(st.ts, tseas, iter_sub, ts_end_)
st = step_fn(fluxes, st, intrpg, cap_s)
no_progress = st.active & (st.tt == tt_before)
st = st._replace(
reason=jnp.where(no_progress, REASON_PROGRESS, st.reason),
active=st.active & ~no_progress,
)
st = _snap_to_substep(st, tt_before, ts_next, cap_s, tseas, iter_sub)
# Pace: a whole window per _pace_denom sub-cycles, from ts_start_.
n_done = i + 1
behind = (
st.active
& (n_done >= _N_WARMUP)
& (
st.ts
< ts_start_ + jnp.float32(n_done - _N_WARMUP) / jnp.float32(_pace_denom)
)
)
st = st._replace(
reason=jnp.where(behind, REASON_SLOW, st.reason),
active=st.active & ~behind,
)
return st, i + 1, ts_end_, ts_start_
state, n_sub, _, _ = lax.while_loop(
cond_fn, body_fn, (state, 0, ts_end, ts_start)
)
stuck = state.active & (state.ts < ts_end) & (n_sub >= max_sub_iter)
n_stuck = jnp.sum(stuck)
state = state._replace(
reason=jnp.where(stuck, REASON_ITERLIMIT, state.reason),
active=state.active & ~stuck,
)
return state, n_sub, n_stuck
def _integrate_window(
state: ParticleState,
step_fn,
fluxes,
*,
tseas: float,
iter_sub: int,
max_sub_iter: int,
timax_seconds: float,
ts_window_end: float,
pace_factor: int = 1,
ts_window_start: float | None = None,
) -> tuple[ParticleState, int]:
"""Advance particles within a single GCM time window [ts_start, ts_start+1).
Runs sub-cycles until either all particles are inactive, every active
particle's ts has reached ts_window_end, or max_sub_iter is reached.
The time interpolation weight intrpg = frac(ts) is updated per sub-cycle,
matching the Fortran mod_clock.F90 :: update_time() logic:
intrpg = MOD(ts, 1.0)
intrpr = 1.0 - intrpg
*step_fn* must have the signature
``step_fn(fluxes, state, intrpg, dsmin) -> ParticleState``.
Delegates the while-loop to :func:`_run_window`, a module-level
``jax.jit``-compiled function. ``fluxes`` is passed as a *dynamic*
JIT argument so JAX traces it by abstract type only — the same compiled
XLA program is reused for every window regardless of the concrete flux
data, preventing per-call CUBIN accumulation in GPU VRAM.
Parameters
----------
state : ParticleState
step_fn : callable
fluxes : FluxFields or LocalFluxFields
tseas : float
iter_sub : int
max_sub_iter : int
timax_seconds : float
ts_window_end : float
pace_factor : int, optional
Divisor applied to the progressive-deactivation denominator.
``1`` (default) keeps the original ``max_sub_iter``-based pace.
``5`` makes the pace 5× stricter: any particle that would need more
than ``max_sub_iter / 5`` sub-cycles to complete the window is
deactivated early. Use ``5`` for multi-window runs where convergence-
zone oscillators inflate the sub-cycle count.
ts_window_start : float or None, optional
Where this call starts, for the slow-particle pace check. Default
``ts_window_end - 1`` (a whole window). Set it when integrating
part of a window (ramp-up saves).
Returns
-------
tuple[ParticleState, int]
Updated particle state and the number of sub-cycles executed.
If the sub-cycle limit is reached, particles still below
``ts_window_end`` are deactivated and a ``RuntimeWarning`` is issued.
"""
# Transfer flux arrays to device here so their GPU lifetime is scoped to
# this function call. LazyFluxFields.compute() returns numpy arrays;
# device_put transfers them to GPU. On return the local `fluxes` goes out
# of scope and the GPU buffers are freed immediately.
fluxes = jax.device_put(fluxes)
ts_end = jnp.asarray(ts_window_end, dtype=jnp.float32)
if ts_window_start is None:
ts_window_start = ts_window_end - 1.0
ts_start = jnp.asarray(ts_window_start, dtype=jnp.float32)
state, n_sub, n_stuck = _run_window(
state,
fluxes,
ts_end,
ts_start,
step_fn=step_fn,
tseas=tseas,
iter_sub=iter_sub,
max_sub_iter=max_sub_iter,
timax_seconds=timax_seconds,
pace_factor=pace_factor,
)
# Single host sync for reporting / warning.
n_sub_val = int(n_sub)
if n_sub_val >= max_sub_iter:
n_stuck_val = int(n_stuck)
if n_stuck_val > 0:
warnings.warn(
f"Sub-cycle limit ({max_sub_iter}) reached; {n_stuck_val} particle(s) "
f"deactivated (ts < {float(ts_end):.1f}). Consider a shorter "
"timestep or increasing max_sub_iter_per_window.",
RuntimeWarning,
stacklevel=3,
)
return state, n_sub_val
def _resume_paused(state, grid, ts_end: float, timax_seconds: float, kmt_np=None):
"""Reactivate particles paused at *ts_end* (a window end or ramp-up save).
A paused particle is inactive with ``reason == 0``, has reached
``ts_end`` (which excludes particles stopped mid-window), is within its
lifetime and sits in a valid ocean cell. With *kmt_np* (numpy ``kmt``)
the work is done in numpy and *state* stays on the host (the batched
path); otherwise in JAX.
"""
xp, kmt = (np, kmt_np) if kmt_np is not None else (jnp, grid.kmt)
# Same float32 value the step cap snapped ts to, so ">=" is exact.
ts_end = np.float32(ts_end)
i0 = xp.clip(state.ia - 1, 0, grid.imt - 1)
j0 = xp.clip(state.ja - 1, 0, grid.jmt - 1)
return state._replace(
active=state.active
| (
~state.active
& (state.reason == 0)
& (state.tt <= state.t0 + timax_seconds)
& (state.ts >= ts_end)
& (kmt[i0, j0] > 0)
& (state.ja >= 1)
& (state.ja <= grid.jmt)
& (state.ka >= 1)
& (state.ka <= grid.km)
)
)
def _seed_windows(seed_freq, grid, fluxes, n_windows: int) -> list[int]:
"""Repeat-seed windows of a run starting at ``fluxes.dtm``."""
if not seed_freq:
return []
dtm = getattr(fluxes, "dtm", None)
if dtm is None:
raise ValueError(
"seed_freq needs the run's start time; pass a flux buffer from "
"tracmass.new_flux (it carries .dtm)"
)
return seed_windows(seed_freq, dtm, grid.timestep, n_windows, grid.reverse)
def _timax(timax_seconds, traj_max_time, grid, fluxes) -> tuple[float, dict]:
"""Lifetime in seconds (and run metadata) of a run starting at ``fluxes.dtm``."""
return resolve_timax(
timax_seconds, traj_max_time, getattr(fluxes, "dtm", None), grid.reverse
)
def _save_time(dtm, grid, frac: float):
"""Time *frac* of a window after the window starting at *dtm*.
Backward runs go back in time, so the save is *frac* windows earlier.
"""
import pandas as pd
step = pd.Timedelta(grid.timestep) * frac
return pd.Timestamp(dtm) - step if grid.reverse else pd.Timestamp(dtm) + step
# ===========================================================================
# Full integration with GCM time stepping
# ===========================================================================
[docs]
def steady_state(
state: ParticleState,
grid: gcm.GridFields,
fluxes: gcm.FluxFields,
*,
iter_sub: int = 1,
max_iter: int = 100_000,
timax_seconds: float = _UNSET,
traj_max_time: str | None = _UNSET,
scheme: str = "tstep",
no_wflux: bool = _UNSET,
n_particle_batches: int = _UNSET,
) -> ParticleState:
"""Single-window steady-state integration.
The flux direction (forward/backward) is determined by ``grid.reverse``,
as set when the grid was created via :func:`tracmass.gcm.new_grid` or
:func:`tracmass.gcm.new_flux`.
Parameters
----------
state : ParticleState
Initial particle state.
grid : gcm.GridFields
Static grid geometry from :func:`tracmass.gcm.new_grid` or
:func:`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
:func:`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 :func:`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
:func:`tracmass.config.n_particle_batches` and :func:`time_dependent`
for the full rationale.
Returns
-------
ParticleState with final positions.
For multi-window integration with field updates, use :func:`time_dependent`.
"""
no_wflux = resolve_no_wflux(no_wflux)
n_particle_batches = resolve_n_particle_batches(n_particle_batches)
timax_seconds, _ = _timax(timax_seconds, traj_max_time, grid, fluxes)
concrete = _ensure_computed(fluxes, no_wflux)
no_wflux = no_wflux or (concrete.wflux is None)
if no_wflux and concrete.wflux is not None:
concrete = concrete._replace(wflux=None)
if scheme == "analytical":
_step = advection_scheme.analytical(
grid, grid.timestep.total_seconds(), no_wflux=no_wflux
)
iter_sub_ = 1
else:
_step = advection_scheme.piecewise(grid, no_wflux=no_wflux)
iter_sub_ = iter_sub
ts_end = float(max_iter)
if n_particle_batches > 1:
# Same particle-batched path as ``time_dependent``: state and flux are
# kept as numpy and pushed to the device one batch at a time inside
# ``_integrate_window``, so peak GPU scratch is ~1/n_particle_batches.
state_np, _ = _integrate_batched(
jax.tree_util.tree_map(np.asarray, state),
_step,
concrete,
n_batches=n_particle_batches,
tseas=grid.timestep.total_seconds(),
iter_sub=iter_sub_,
max_sub_iter=max_iter,
timax_seconds=timax_seconds,
ts_window_end=ts_end,
pace_factor=1,
)
return jax.device_put(state_np)
state, _ = _integrate_window(
state,
_step,
concrete,
tseas=grid.timestep.total_seconds(),
iter_sub=iter_sub_,
max_sub_iter=max_iter,
timax_seconds=timax_seconds,
ts_window_end=ts_end,
)
return state
[docs]
@manage_run_io
def time_dependent(
state: ParticleState,
grid: gcm.GridFields,
fluxes: gcm.FluxFields,
fluxes_fn: Callable | None = None,
*,
iter_sub: int = 1,
max_iter: int = 100_000,
n_windows: int | None = None,
max_sub_iter_per_window: int = 10_000,
timax_seconds: float = _UNSET,
traj_max_time: str | None = _UNSET,
recorder: ParticleRecorder | str | None = _UNSET,
monitor: RunMonitor | str | None = _UNSET,
scheme: str = "tstep",
no_wflux: bool = _UNSET,
pace_factor: int = 5,
n_particle_batches: int = _UNSET,
rampup_windows: int = _UNSET,
rampup_substeps: int = _UNSET,
seed_freq: str | None = _UNSET,
) -> ParticleState:
"""Multi-window integration using the ned (new/update grid) system.
Each window the flux buffer is advanced by one timestep via
:func:`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 :func:`tracmass.gcm.new_grid`
or :func:`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 :func:`~tracmass.gcm.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
:func:`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
:func:`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 :class:`~tracmass.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 :mod:`tracmass.runsetup`.
monitor : RunMonitor or str or None, optional
Controls progress output. Pass a :class:`~tracmass.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 :func:`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 :func:`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
:func:`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``).
Returns
-------
ParticleState with final positions.
"""
tseas = grid.timestep.total_seconds()
no_wflux = resolve_no_wflux(no_wflux)
n_windows = resolve_n_windows(n_windows, fluxes)
n_particle_batches = resolve_n_particle_batches(n_particle_batches)
rampup_windows, rampup_substeps = resolve_rampup(rampup_windows, rampup_substeps)
seed_freq = resolve_seed_freq(seed_freq)
reseed_wins = _seed_windows(seed_freq, grid, fluxes, n_windows)
timax_seconds, timax_meta = _timax(timax_seconds, traj_max_time, grid, fluxes)
initial_state = state
if rampup_windows and scheme == "analytical":
raise ValueError(
"time_dependent: ramp-up saves need scheme='tstep'; the analytical "
"scheme cannot stop particles inside a window"
)
if recorder is not None:
recorder.metadata.update(
rampup_metadata(n_windows, tseas, rampup_windows, rampup_substeps)
)
recorder.metadata.update(seed_metadata(seed_freq, reseed_wins))
recorder.metadata.update(timax_meta)
# Infer fluxes_fn from the lazy object when not supplied explicitly.
if fluxes_fn is None:
fluxes_fn = getattr(fluxes, "fluxes_fn", None)
log.debug(
"time_dependent: %d particles, n_windows=%d, timestep=%s, "
"no_wflux=%s, scheme=%s, iter_sub=%d, n_particle_batches=%d, "
"monitor=%s, recorder=%s",
int(np.asarray(state.active).shape[0]), n_windows, grid.timestep,
no_wflux, scheme, iter_sub, n_particle_batches,
getattr(monitor, "_backend", monitor),
getattr(recorder, "_filepath", recorder),
)
# Peek at no_wflux before the loop; we need the first concrete slice for this.
with timed(log, "time_dependent: loading first window (before monitor starts)"):
first_concrete = _ensure_computed(fluxes, no_wflux)
no_wflux = no_wflux or (first_concrete.wflux is None)
if no_wflux and first_concrete.wflux is not None:
first_concrete = first_concrete._replace(wflux=None)
if scheme == "analytical":
_step = advection_scheme.analytical(grid, tseas, no_wflux=no_wflux)
iter_sub_ = 1
else:
_step = advection_scheme.piecewise(grid, no_wflux=no_wflux)
iter_sub_ = iter_sub
# Use the already-computed first window rather than computing again.
_pending_concrete: FluxFields | None = first_concrete
# Batched path: keep state as numpy (CPU) between windows so the GPU holds
# zero bytes outside of _integrate_window calls. The only GPU memory
# during a window is: one batch state + flux + XLA scratch, all freed when
# _integrate_window returns. Cache kmt on CPU for numpy reactivation.
kmt_cpu: np.ndarray | None = None
if n_particle_batches > 1:
kmt_cpu = np.asarray(grid.kmt)
state = jax.tree_util.tree_map(np.asarray, state)
if monitor is not None:
log.debug("time_dependent: starting monitor")
monitor.on_run_start(
n_windows=n_windows,
n_particles=int(np.asarray(state.active).shape[0]),
)
_last_win = -1
for win in range(n_windows):
_t_win_start = time.perf_counter()
if win in reseed_wins:
state = reseed(state, initial_state, win, tseas)
active_before = np.asarray(state.active).copy()
n_active_before = int(active_before.sum())
if n_active_before == 0:
if not any(w > win for w in reseed_wins):
break
# Nothing to advect until the next release: just step the fluxes.
_pending_concrete = None
if fluxes_fn is not None and win < n_windows - 1:
fluxes = gcm.update_flux(grid, fluxes, fluxes_fn)
_last_win = win
continue
ts_window_end = float(win + 1)
log.debug(
"window %d/%d dtm=%s: %d active", win + 1, n_windows, fluxes.dtm,
n_active_before,
)
if recorder is not None:
with timed(log, "window %d: recording positions", win + 1):
recorder.record(state, dtm=fluxes.dtm, window=win)
if monitor is not None:
monitor.on_window_start(win, n_active_before, fluxes.dtm)
# --- I/O phase: compute fluxes (Dask read for GLORYS, no-op otherwise) ---
_t0 = time.perf_counter()
if _pending_concrete is not None:
concrete = _pending_concrete
_pending_concrete = None
else:
concrete = _ensure_computed(fluxes, no_wflux)
if no_wflux and concrete.wflux is not None:
concrete = concrete._replace(wflux=None)
_t_io = time.perf_counter() - _t0
# --- JAX phase: particle advection (includes device_put + while_loop) ---
# One piece per save time: normally the whole window; during ramp-up
# rampup_substeps pieces, saving positions between them.
_t0 = time.perf_counter()
ends = window_ends(win, rampup_windows, rampup_substeps)
n_sub = 0
ts_start = float(win)
for i_end, ts_end in enumerate(ends):
if n_particle_batches > 1:
# state is already numpy; _integrate_batched returns numpy.
# GPU holds nothing between windows — only one batch's state +
# flux + XLA scratch are live on device at any one time, and
# all three are freed when _integrate_window returns.
state, n_piece = _integrate_batched(
state,
_step,
concrete,
n_batches=n_particle_batches,
tseas=tseas,
iter_sub=iter_sub_,
max_sub_iter=max_sub_iter_per_window,
timax_seconds=timax_seconds,
ts_window_end=ts_end,
pace_factor=pace_factor,
ts_window_start=ts_start,
)
else:
state, n_piece = _integrate_window(
state,
_step,
concrete,
tseas=tseas,
iter_sub=iter_sub_,
max_sub_iter=max_sub_iter_per_window,
timax_seconds=timax_seconds,
ts_window_end=ts_end,
pace_factor=pace_factor,
ts_window_start=ts_start,
)
n_sub += n_piece
if i_end < len(ends) - 1:
state = _resume_paused(state, grid, ts_end, timax_seconds, kmt_cpu)
if recorder is not None:
recorder.record_rampup(
state, dtm=_save_time(fluxes.dtm, grid, (i_end + 1) / len(ends)),
window=win,
)
ts_start = ts_end
_t_jax = time.perf_counter() - _t0
log.debug(
"window %d: flux I/O %.1f s, advection %.1f s%s, %d sub-cycles",
win + 1, _t_io, _t_jax, " (incl. JIT compile)" if win == 0 else "",
n_sub,
)
# Tally permanent deactivations in this window (reason != 0 means
# permanently stopped, not just paused at the window boundary).
reasons_arr = np.asarray(state.reason)
newly_stopped = active_before & (reasons_arr != 0)
n_stopped = int(newly_stopped.sum())
reason_tally: dict[str, int] = {}
if n_stopped > 0:
for code in sorted(_REASON_NAMES):
if code == 0:
continue
cnt = int(np.sum(newly_stopped & (reasons_arr == code)))
if cnt:
reason_tally[_REASON_NAMES[code]] = cnt
if recorder is not None and n_stopped > 0:
recorder.record_deactivations(
state, newly_stopped, dtm=fluxes.dtm, window=win
)
# Drop the reference and force GC so JAX can reclaim device memory
# before the next window's flux data is loaded.
del concrete
gc.collect()
# Reactivate particles paused at the window boundary.
# Batched path: numpy, so no GPU memory is consumed between windows.
state = _resume_paused(state, grid, ts_window_end, timax_seconds, kmt_cpu)
# Advance the flux buffer to the next time window (zero-cost for LazyFluxFields).
if fluxes_fn is not None and win < n_windows - 1:
with timed(log, "window %d: update_flux", win + 1):
fluxes = gcm.update_flux(grid, fluxes, fluxes_fn)
_t_win = time.perf_counter() - _t_win_start
_t_py = max(0.0, _t_win - _t_io - _t_jax)
n_active_after = int(np.sum(np.asarray(state.active)))
if monitor is not None:
monitor.on_window_end(
win=win,
dtm=fluxes.dtm,
n_active_before=n_active_before,
n_active_after=n_active_after,
n_stopped=n_stopped,
reason_tally=reason_tally,
n_sub=n_sub,
limit_hit=n_sub >= max_sub_iter_per_window,
t_io=_t_io,
t_jax=_t_jax,
t_py=_t_py,
)
_last_win = win
if monitor is not None:
monitor.on_run_end()
if recorder is not None:
recorder.record(state, dtm=fluxes.dtm, window=_last_win + 1)
# Batched path kept state as numpy throughout; put it back on device before
# returning so callers always receive a GPU ParticleState.
if n_particle_batches > 1:
state = jax.device_put(state)
return state
