"""Memory-efficient integration variants for TRACMASS-JAX.
Provides two integration functions that avoid loading the full global flux
grid into device memory at once:
- :func:`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.
- :func:`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 :func:`~tracmass.integrate.full_grid._integrate_window`
internally (patched variant) or replace it with an explicit Python loop
(streaming variant).
"""
from __future__ import annotations
import gc
import time
import jax
from typing import Callable
import jax.numpy as jnp
import numpy as np
from tracmass import advection_scheme, gcm
from tracmass.gcm.flux_grid import FluxStencil, extract_stencil
from tracmass.integrate.full_grid import (
_integrate_window,
_resume_paused,
_save_time,
_seed_windows,
_timax,
_snap_to_substep,
_substep_cap,
)
from tracmass.monitoring import RunMonitor
from tracmass.particle import (
_REASON_NAMES,
REASON_TIME,
ParticleRecorder,
ParticleState,
)
from tracmass.runsetup import (
_UNSET,
manage_run_io,
resolve_n_particle_batches,
resolve_n_windows,
resolve_no_wflux,
resolve_rampup,
resolve_seed_freq,
rampup_metadata,
reseed,
seed_metadata,
window_ends,
)
# ---------------------------------------------------------------------------
# Stencil-cache helper
# ---------------------------------------------------------------------------
def _update_stencil(
cached: FluxStencil,
delta: FluxStencil,
moved: jnp.ndarray,
) -> FluxStencil:
"""Return *cached* with *delta* values substituted for *moved* particles.
Parameters
----------
cached : FluxStencil
Stencil values from the previous sub-cycle.
delta : FluxStencil
Freshly extracted stencil. Only positions where *moved* is ``True``
carry valid data; all other positions are zeros (as produced by
:func:`~tracmass.gcm.extract_stencil` with a mask argument).
moved : jnp.ndarray
Boolean mask, shape ``(N,)``. ``True`` for particles that crossed
a cell wall since the last sub-cycle.
Returns
-------
FluxStencil
Merged stencil: *delta* values where *moved*, *cached* elsewhere.
"""
sel2 = moved[:, None] # (N, 1) — broadcast over the 2-slot axis
def _m2(c: jnp.ndarray, d: jnp.ndarray) -> jnp.ndarray:
return jnp.where(sel2, d, c)
def _m1(c: jnp.ndarray, d: jnp.ndarray) -> jnp.ndarray:
return jnp.where(moved, d, c)
w_top = None if cached.w_top is None else _m2(cached.w_top, delta.w_top)
w_bot = None if cached.w_bot is None else _m2(cached.w_bot, delta.w_bot)
return FluxStencil(
u_east = _m2(cached.u_east, delta.u_east),
u_west = _m2(cached.u_west, delta.u_west),
v_north = _m2(cached.v_north, delta.v_north),
v_south = _m2(cached.v_south, delta.v_south),
w_top = w_top,
w_bot = w_bot,
dzt = _m2(cached.dzt, delta.dzt),
zstot = _m2(cached.zstot, delta.zstot),
kmt = _m1(cached.kmt, delta.kmt),
dxdy = _m1(cached.dxdy, delta.dxdy),
)
# ---------------------------------------------------------------------------
# Patch helpers
# ---------------------------------------------------------------------------
def _round_up(x: int, granularity: int) -> int:
"""Round *x* up to the nearest multiple of *granularity*."""
return ((x + granularity - 1) // granularity) * granularity
def _zarr_chunk_xy(fluxes) -> tuple[int, int]:
"""Return ``(chunk_x, chunk_y)`` of the underlying Zarr store, or ``(1, 1)``.
Reads from :attr:`LazyFluxFields.zarr_chunk_xy`, which is set by GCM
adapters (e.g. ``gcm.glorys``) from the raw velocity DataArray *before*
arithmetic operations that would corrupt the chunk layout. Returns
``(1, 1)`` for in-memory adapters (theoretical, fabboni) where
chunk alignment has no effect.
"""
return getattr(fluxes, "zarr_chunk_xy", (1, 1)) or (1, 1)
def _chunk_bbox(
ia_arr: np.ndarray,
ja_arr: np.ndarray,
buffer: int,
zarr_cx: int,
zarr_cy: int,
imt_grid: int,
jmt_grid: int,
) -> tuple[int, int, int, int]:
"""Return a chunk-aligned ``(pi0, pj0, imt_p, jmt_p)`` covering particles.
Snaps the lower-left corner of the bounding box (including *buffer*)
down to the nearest Zarr chunk boundary and rounds the size up to cover
the upper-right corner. The resulting patch always starts **and ends**
on a chunk boundary, so ``compute_patch`` reads exactly
``(imt_p / zarr_cx) × (jmt_p / zarr_cy)`` complete chunks — no partial
chunk reads, no wasted decompression.
``ia_arr`` / ``ja_arr`` are **1-based** TRACMASS cell indices.
Parameters
----------
ia_arr, ja_arr : np.ndarray
1-based cell indices of the active particle population.
buffer : int
Extra cells added on each side before snapping to chunk boundaries.
zarr_cx, zarr_cy : int
Zarr chunk sizes in x and y. Pass ``(1, 1)`` to get the same
result as the old unaligned logic.
imt_grid, jmt_grid : int
Full-grid dimensions used for clamping.
Returns
-------
tuple[int, int, int, int]
``(pi0, pj0, imt_p, jmt_p)`` — all 0-based, sizes exclusive-end.
"""
# 0-based lower and upper bounds including buffer
i_lo = max(0, int(ia_arr.min()) - 1 - buffer)
i_hi = min(imt_grid - 1, int(ia_arr.max()) - 1 + buffer)
j_lo = max(0, int(ja_arr.min()) - 1 - buffer)
j_hi = min(jmt_grid - 1, int(ja_arr.max()) - 1 + buffer)
# Snap to chunk boundaries
pi0 = (i_lo // zarr_cx) * zarr_cx
pj0 = (j_lo // zarr_cy) * zarr_cy
imt_p = min(_round_up(i_hi - pi0 + 1, zarr_cx), imt_grid - pi0)
jmt_p = min(_round_up(j_hi - pj0 + 1, zarr_cy), jmt_grid - pj0)
return pi0, pj0, imt_p, jmt_p
def _shift_state(
state: ParticleState, di: int, dj: int
) -> ParticleState:
"""Shift particle grid indices and fractional positions by ``(di, dj)``.
Used to convert between global coordinates and patch-local coordinates.
``di`` and ``dj`` are added to ``ia``, ``ja``, ``x``, ``y``.
Parameters
----------
state : ParticleState
di, dj : int
Offset to add (use negative values to convert global → local).
"""
return state._replace(
ia=state.ia + di,
ja=state.ja + dj,
x=state.x + di,
y=state.y + dj,
)
# ===========================================================================
# Patch-based multi-window integration
# ===========================================================================
[docs]
@manage_run_io
def time_dependent_patched(
state: ParticleState,
grid: gcm.GridFields,
fluxes: gcm.LazyFluxFields,
fluxes_fn=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,
no_wflux: bool = _UNSET,
patch_buffer: int = 50,
patch_granularity: int = 100,
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 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
:func:`~tracmass.gcm.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
:func:`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
: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`.
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
:func:`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 :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.
n_particle_batches : int, optional
Accepted for API compatibility with the full-grid integrators
(:func:`~tracmass.integrate.steady_state` /
:func:`~tracmass.integrate.time_dependent`) — resolved from the
``n_particle_batches`` setting when omitted — but has no effect here.
Default ``1``.
rampup_windows, rampup_substeps : int, optional
Extra ``event_type="rampup"`` saves in the first windows, as in
:func:`~tracmass.integrate.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
:func:`~tracmass.integrate.time_dependent`. When omitted, from
``seed_freq`` in the ``[env.seed]`` settings table.
Returns
-------
ParticleState with final positions in **global** grid coordinates.
"""
tseas = grid.timestep.total_seconds()
no_wflux = resolve_no_wflux(no_wflux)
n_windows = resolve_n_windows(n_windows, fluxes)
# Resolved for API symmetry; the streaming loop does not batch particles.
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 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)
if fluxes_fn is None:
fluxes_fn = getattr(fluxes, "fluxes_fn", None)
# ------------------------------------------------------------------
# 1. Initial patch shape — grows as the particle cloud spreads
# ------------------------------------------------------------------
# Detect Zarr chunk sizes so patches are chunk-aligned: each compute_patch
# call reads complete chunks with no cross-boundary waste. Falls back to
# (1, 1) for in-memory adapters (theoretical, fabboni) where alignment
# has no effect.
zarr_cx, zarr_cy = _zarr_chunk_xy(fluxes)
active_mask_init = np.asarray(state.active)
ia_init = np.asarray(state.ia)[active_mask_init]
ja_init = np.asarray(state.ja)[active_mask_init]
pi0, pj0, imt_p, jmt_p = _chunk_bbox(
ia_init, ja_init, patch_buffer, zarr_cx, zarr_cy, grid.imt, grid.jmt
)
# Compile the local kernel for the initial patch shape.
# piecewise_local is LRU-cached, so later calls with the same
# (imt_p, jmt_p) return the already-compiled kernel without retracing.
_step_local = advection_scheme.piecewise_local(
imt_p, jmt_p, grid.km, grid.iperio, no_wflux=no_wflux
)
if monitor is not None:
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)
if not np.any(np.asarray(state.active)):
if not any(w > win for w in reseed_wins):
break
# Nothing to advect until the next release: just step the fluxes.
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)
active_before = np.asarray(state.active).copy()
n_active_before = int(active_before.sum())
if recorder is not None:
recorder.record(state, dtm=fluxes.dtm, window=win)
# ------------------------------------------------------------------
# 2. Resize (if needed) and re-centre patch on particle bounding box
# ------------------------------------------------------------------
active_mask = np.asarray(state.active)
ia_arr = np.asarray(state.ia)[active_mask]
ja_arr = np.asarray(state.ja)[active_mask]
# Recompute chunk-aligned origin and minimum size for this window.
# Never shrink imt_p/jmt_p: avoids oscillating recompilations.
pi0_new, pj0_new, imt_p_new, jmt_p_new = _chunk_bbox(
ia_arr, ja_arr, patch_buffer, zarr_cx, zarr_cy, grid.imt, grid.jmt
)
new_imt_p = min(max(imt_p, imt_p_new), grid.imt)
new_jmt_p = min(max(jmt_p, jmt_p_new), grid.jmt)
if new_imt_p != imt_p or new_jmt_p != jmt_p:
imt_p, jmt_p = new_imt_p, new_jmt_p
_step_local = advection_scheme.piecewise_local(
imt_p, jmt_p, grid.km, grid.iperio, no_wflux=no_wflux
)
# Use the chunk-aligned origin; clamp so the (possibly larger)
# patch doesn't exceed the grid boundary.
pi0 = int(np.clip(pi0_new, 0, grid.imt - imt_p))
pj0 = int(np.clip(pj0_new, 0, grid.jmt - jmt_p))
if monitor is not None:
monitor.on_window_start(win, n_active_before, fluxes.dtm)
# ------------------------------------------------------------------
# 3. Compute only the patch (Dask → JAX)
# ------------------------------------------------------------------
_t0 = time.perf_counter()
local_flux = fluxes.compute_patch(
pi0, pi0 + imt_p, pj0, pj0 + jmt_p, grid,
include_wflux=not no_wflux,
)
_t_io = time.perf_counter() - _t0
# ------------------------------------------------------------------
# 4. Convert particle coords to patch-local and run the window
# ------------------------------------------------------------------
# One piece per save time (several during ramp-up); positions are
# saved and paused particles resumed in global coords 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):
state_local, n_piece = _integrate_window(
_shift_state(state, -pi0, -pj0),
_step_local,
local_flux,
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
# Convert back to global coords.
state = _shift_state(state_local, pi0, pj0)
if i_end < len(ends) - 1:
state = _resume_paused(state, grid, ts_end, timax_seconds)
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
del local_flux
gc.collect()
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 newly_stopped.any():
recorder.record_deactivations(
state, newly_stopped, dtm=fluxes.dtm, window=win
)
# Reactivate particles paused at the window boundary.
# Require ts >= ts_window_end to exclude particles deactivated
# mid-window (stuck, on-land, or out-of-bounds).
tt_np = np.asarray(state.tt)
t0_np = np.asarray(state.t0)
ia_p = np.clip(np.asarray(state.ia) - 1, 0, grid.imt - 1)
ja_p = np.clip(np.asarray(state.ja) - 1, 0, grid.jmt - 1)
still_valid = (
(np.asarray(state.reason) == 0)
& (np.asarray(grid.kmt)[ia_p, ja_p] > 0)
& (np.asarray(state.ja) >= 1)
& (np.asarray(state.ja) <= grid.jmt)
& (np.asarray(state.ka) >= 1)
& (np.asarray(state.ka) <= grid.km)
)
state = state._replace(
active=jnp.asarray(
np.asarray(state.active)
| (
~np.asarray(state.active)
& ((tt_np - t0_np) <= timax_seconds)
& (np.asarray(state.ts) >= ts_window_end)
& still_valid
)
)
)
# Advance lazy flux buffer to next window (zero-cost reference swap).
if fluxes_fn is not None and win < n_windows - 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)
return state
# ===========================================================================
# Streaming (per-particle stencil) multi-window integration
# ===========================================================================
[docs]
@manage_run_io
def time_dependent_streaming(
state: ParticleState,
grid: gcm.GridFields,
fluxes: gcm.LazyFluxFields,
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,
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 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 :func:`~tracmass.integrate.full_grid.time_dependent`
and :func:`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
:func:`~tracmass.gcm.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
: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.
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.
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 :func:`tracmass.config.move_vertically`.
n_particle_batches : int, optional
Accepted for API compatibility with the full-grid integrators
(:func:`~tracmass.integrate.steady_state` /
:func:`~tracmass.integrate.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, rampup_substeps : int, optional
Not supported here: a ramp-up (from the arguments or the
``rampup_windows`` setting) raises ``NotImplementedError``; use
:func:`~tracmass.integrate.time_dependent` or
:func:`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
:func:`~tracmass.integrate.time_dependent`. When omitted, from
``seed_freq`` in the ``[env.seed]`` settings table.
Returns
-------
ParticleState
Final particle state in global grid coordinates.
"""
tseas = grid.timestep.total_seconds()
no_wflux = resolve_no_wflux(no_wflux)
n_windows = resolve_n_windows(n_windows, fluxes)
# Resolved for API symmetry; the streaming loop does not batch particles.
n_particle_batches = resolve_n_particle_batches(n_particle_batches)
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 recorder is not None:
recorder.metadata.update(seed_metadata(seed_freq, reseed_wins))
recorder.metadata.update(timax_meta)
if fluxes_fn is None:
fluxes_fn = getattr(fluxes, "fluxes_fn", None)
if resolve_rampup(rampup_windows, rampup_substeps)[0]:
raise NotImplementedError(
"time_dependent_streaming does not support ramp-up saves "
"(rampup_windows); use time_dependent or time_dependent_patched, "
"or set rampup_windows = 0"
)
has_w = (
fluxes._past.get("wflux") is not None
and fluxes._present.get("wflux") is not None
)
no_wflux = no_wflux or not has_w
# Compile kernel once — closes over grid (kmt for land checks)
_step = advection_scheme.piecewise_stencil(grid, no_wflux=no_wflux)
if monitor is not None:
monitor.on_run_start(
n_windows=n_windows,
n_particles=int(np.asarray(state.active).shape[0]),
)
total_iters = 0
_last_win = -1
for win in range(n_windows):
_t_win_start = time.perf_counter()
ts_window_start = float(win)
ts_window_end = float(win + 1)
if win in reseed_wins:
state = reseed(state, initial_state, win, tseas)
# t0 is the per-particle trajectory start time; it never changes during
# the inner sub-cycle loop, so we fetch it once per window.
t0_np = np.asarray(state.t0)
if not np.any(np.asarray(state.active)):
if not any(w > win for w in reseed_wins):
break
# Nothing to advect until the next release: just step the fluxes.
if fluxes_fn is not None and win < n_windows - 1:
fluxes = gcm.update_flux(grid, fluxes, fluxes_fn)
_last_win = win
continue
active_before = np.asarray(state.active).copy()
n_active_before = int(active_before.sum())
if recorder is not None:
recorder.record(state, dtm=fluxes.dtm, window=win)
if monitor is not None:
monitor.on_window_start(win, n_active_before, fluxes.dtm)
_t_io_win = 0.0
_t_jax_win = 0.0
win_sub = 0
# Preload all Dask arrays for this window into numpy in one batch
# compute, so that extract_stencil calls inside the sub-cycle loop
# use pure numpy indexing with no further Zarr I/O.
_t0 = time.perf_counter()
_fluxes = fluxes.preload(include_wflux=not no_wflux)
_t_io_win += time.perf_counter() - _t0
# Progressive deactivation denominator: same pace_factor=5 as the
# while_loop variant. Particles that would need more than
# max_sub_iter_per_window / 5 steps are detected as stuck and killed.
_N_WARMUP = 2
_pace_denom = max((max_sub_iter_per_window - _N_WARMUP) // pace_factor, 1)
# Stencil cache: reset each window because the flux data changed.
# Within a window, the stencil is only reloaded for particles that
# cross a cell wall — particles staying in the same cell reuse the
# previously extracted values, saving Dask I/O.
stencil_cache: FluxStencil | None = None
prev_ia = prev_ja = prev_ka = None
for it in range(max_sub_iter_per_window):
if total_iters >= max_iter:
break
# One batched device→host transfer per sub-cycle iteration.
# ia/ja/ka are fetched here to detect cell crossings for caching.
active_np, tt_np, ia_np, ja_np, ka_np = jax.device_get(
(state.active, state.tt, state.ia, state.ja, state.ka)
)
if not np.any(active_np):
break
ts_np = tt_np / tseas
# Deactivate particles that exceeded lifetime or passed window end
# t0_np was fetched once before the inner loop (t0 is constant)
exceeded = (tt_np - t0_np) > timax_seconds
past_window = ts_np >= ts_window_end
new_active = active_np & ~exceeded & ~past_window
if not np.array_equal(active_np, new_active):
state = state._replace(
active=jnp.asarray(new_active),
reason=jnp.where(
jnp.asarray(active_np & exceeded), REASON_TIME, state.reason
),
)
active_np = new_active
# Progressive deactivation: by step n a particle must have covered
# at least (n - _N_WARMUP) / _pace_denom of the window. Kills
# convergence-zone oscillators before they inflate the loop count.
if it >= _N_WARMUP:
pace_threshold = (ts_window_start
+ (it - _N_WARMUP) / _pace_denom)
too_slow = active_np & (ts_np < pace_threshold)
if np.any(too_slow):
new_active = active_np & ~too_slow
state = state._replace(active=jnp.asarray(new_active))
active_np = new_active
if not np.any(active_np):
break
# Check if all active particles have exited this window
if np.all(past_window[active_np]):
break
# Time-interpolation weight from mean active-particle ts
intrpg = float(np.mean(ts_np[active_np]) % 1.0)
# Stencil cache: extract only for active particles that are in a
# new cell since the previous sub-cycle. On the first sub-cycle
# of each window (stencil_cache is None) extract for all active
# particles; thereafter only re-read for those that crossed a wall.
if stencil_cache is None:
_t0 = time.perf_counter()
stencil_cache = extract_stencil(
_fluxes, grid, state, no_wflux=no_wflux
)
_t_io_win += time.perf_counter() - _t0
else:
moved = active_np & (
(ia_np != prev_ia)
| (ja_np != prev_ja)
| (ka_np != prev_ka)
)
if np.any(moved):
_t0 = time.perf_counter()
delta = extract_stencil(
_fluxes, grid, state, no_wflux=no_wflux, mask=moved
)
_t_io_win += time.perf_counter() - _t0
stencil_cache = _update_stencil(
stencil_cache, delta, jnp.asarray(moved)
)
# else: no cell crossings — reuse stencil_cache unchanged
prev_ia = ia_np.copy()
prev_ja = ja_np.copy()
prev_ka = ka_np.copy()
# One sub-cycle for all particles
_t0 = time.perf_counter()
tt_before = state.tt
ts_next, cap_s = _substep_cap(state.ts, tseas, iter_sub)
state = _step(stencil_cache, state, intrpg, cap_s)
_t_jax_win += time.perf_counter() - _t0
# Keep ts in sync with tt — pure JAX, no device round-trip
state = _snap_to_substep(
state, tt_before, ts_next, cap_s, tseas, iter_sub
)
total_iters += 1
win_sub += 1
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 newly_stopped.any():
recorder.record_deactivations(
state, newly_stopped, dtm=fluxes.dtm, window=win
)
# Advance lazy flux buffer to next window (zero-cost reference swap)
if fluxes_fn is not None and win < n_windows - 1:
fluxes = gcm.update_flux(grid, fluxes, fluxes_fn)
# Reactivate particles paused at the window boundary.
# Require ts >= ts_window_end to exclude particles deactivated
# mid-window (stuck, on-land, or out-of-bounds).
tt_np = np.asarray(state.tt)
i0_p = np.clip(np.asarray(state.ia) - 1, 0, grid.imt - 1)
j0_p = np.clip(np.asarray(state.ja) - 1, 0, grid.jmt - 1)
still_valid = (
(np.asarray(state.reason) == 0)
& (np.asarray(grid.kmt)[i0_p, j0_p] > 0)
& (np.asarray(state.ja) >= 1)
& (np.asarray(state.ja) <= grid.jmt)
& (np.asarray(state.ka) >= 1)
& (np.asarray(state.ka) <= grid.km)
)
state = state._replace(
active=jnp.asarray(
np.asarray(state.active)
| (
~np.asarray(state.active)
& ((tt_np - t0_np) <= timax_seconds)
& (np.asarray(state.ts) >= ts_window_end)
& still_valid
)
)
)
_t_win = time.perf_counter() - _t_win_start
_t_py = max(0.0, _t_win - _t_io_win - _t_jax_win)
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=win_sub,
limit_hit=win_sub >= max_sub_iter_per_window,
t_io=_t_io_win,
t_jax=_t_jax_win,
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)
return state
