from __future__ import annotations
from typing import NamedTuple
import jax.numpy as jnp
import numpy as np
# ---------------------------------------------------------------------------
# Deactivation reason codes (stored in ParticleState.reason, dtype int16)
# ---------------------------------------------------------------------------
REASON_ACTIVE = jnp.int16(0) # still running / window-paused
REASON_TIME = jnp.int16(1) # time limit exceeded (tt − t0 > timax)
REASON_LAND = jnp.int16(2) # hit land cell (kmt == 0)
REASON_BOUNDS_H = jnp.int16(3) # left model domain horizontally (j out of range)
REASON_PROGRESS = jnp.int16(4) # zero float32 progress (ds ≈ 0 / UNDEF)
REASON_SLOW = jnp.int16(5) # progressive deactivation (too slow)
REASON_ITERLIMIT = jnp.int16(6) # sub-cycle iteration limit reached
REASON_DXYZ = jnp.int16(7) # zero/negative cell volume (dxyz ≤ 0)
REASON_BOUNDS_V = jnp.int16(8) # left model domain vertically (k out of range)
_REASON_NAMES: dict[int, str] = {
0: "active",
1: "time_limit",
2: "land",
3: "out_of_bounds_h",
4: "no_progress",
5: "too_slow",
6: "iter_limit",
7: "zero_volume",
8: "out_of_bounds_v",
}
# ---------------------------------------------------------------------------
# Particle state (Structure-of-Arrays via a NamedTuple)
# ---------------------------------------------------------------------------
[docs]
class ParticleState(NamedTuple):
"""State vector for all particles — each field is shape (n,)."""
x: jnp.ndarray # zonal position (1-based dimensionless)
y: jnp.ndarray # meridional position
z: jnp.ndarray # vertical position
ia: jnp.ndarray # zonal box index (1-based int)
ja: jnp.ndarray # meridional box index
ka: jnp.ndarray # vertical box index
tt: jnp.ndarray # elapsed time [s]
ts: jnp.ndarray # normalised time (tt / tseas)
t0: jnp.ndarray # seeding time [s]
subvol: jnp.ndarray # transported volume
niter: jnp.ndarray # iteration counter
active: jnp.ndarray # bool mask — still advecting?
reason: jnp.ndarray # int16 — why this particle stopped (0 = still active)
[docs]
@classmethod
def concat(cls, *states: ParticleState) -> ParticleState:
"""Join particle states end to end.
Particles keep their order, so the particles of the first state keep
their indices (the ``id`` written by :class:`ParticleRecorder`) and
those of later states get new ones after them.
Parameters
----------
*states : ParticleState
States to join, in order.
Returns
-------
ParticleState
NumPy arrays if every input field is a NumPy array (e.g. the
host-side state of a batched run), JAX arrays otherwise.
Examples
--------
>>> both = ParticleState.concat(state, new_particles)
>>> len(both.x) == len(state.x) + len(new_particles.x)
True
"""
if not states:
raise ValueError("ParticleState.concat needs at least one state")
xp = np if all(isinstance(a, np.ndarray) for s in states for a in s) else jnp
return cls(*(xp.concatenate(fields) for fields in zip(*states)))
def __repr__(self):
lines = ["ParticleState("]
lines.append(f"Total particles: {len(self.x)}")
n_active = int(np.sum(self.active))
lines.append(f"Active particles: {n_active}")
n_inactive = len(self.x) - n_active
if n_inactive > 0:
reason_arr = np.asarray(self.reason)
inactive_mask = ~np.asarray(self.active)
parts = []
for code, name in _REASON_NAMES.items():
if code == 0:
continue
cnt = int(np.sum(inactive_mask & (reason_arr == code)))
if cnt:
parts.append(f"{name}:{cnt}")
if parts:
lines.append(f"Stopped ({n_inactive}): " + ", ".join(parts))
for name, value in zip(self._fields, self):
if hasattr(value, "shape"):
lines.append(f" {name}: shape={value.shape}, dtype={value.dtype},")
else:
lines.append(f" {name}: {value},")
lines.append(")")
return "\n".join(lines)
[docs]
def write_trajectories_parquet(
filepath: str,
state: ParticleState,
nff: int = 1,
) -> str:
"""Write a single particle state snapshot to a Parquet file.
Parameters
----------
filepath : str
Output ``.parquet`` file path.
state : ParticleState
Final (or snapshot) particle state.
nff : int
Time direction (1=forward, -1=backward). ``tt`` and ``ts`` are
stored as ``nff * value`` so backward trajectories have negative
times, matching the Fortran TRACMASS convention.
Returns
-------
str
The path that was written.
"""
import pandas as pd
n = len(state.x)
df = pd.DataFrame({
"ntrac": np.arange(n, dtype=np.int32),
"x": np.asarray(state.x),
"y": np.asarray(state.y),
"z": np.asarray(state.z),
"ia": np.asarray(state.ia),
"ja": np.asarray(state.ja),
"ka": np.asarray(state.ka),
"tt": np.asarray(state.tt) * nff,
"ts": np.asarray(state.ts) * nff,
"t0": np.asarray(state.t0),
"subvol": np.asarray(state.subvol),
"niter": np.asarray(state.niter),
"active": np.asarray(state.active),
"reason": np.asarray(state.reason),
})
df.to_parquet(filepath, index=False)
return filepath
[docs]
class ParticleRecorder:
"""Streams trajectory snapshots and deactivation events to a Parquet file.
Each call to :meth:`record` appends active-particle positions as a new
row group with ``event_type="snapshot"`` (at a window boundary); each
call to :meth:`record_rampup` the same with ``event_type="rampup"``
(extra saves inside the first windows of a run). Each call to
:meth:`record_deactivations` appends newly-stopped particles with
``event_type="deactivated"``. Call :meth:`close` (or use the recorder
as a context manager) to finalise the file.
The unified schema contains 13 columns: ``event_type``, ``window``,
``dtm``, ``id``, ``x``, ``y``, ``z``, ``ia``, ``ja``, ``ka``, ``tt``,
``t0``, ``reason``. Filter on ``event_type`` when analysing a single
event class.
``x`` / ``ia`` and ``y`` / ``ja`` are written in full-grid (absolute)
indices: :attr:`grid_offset` — set from ``grid.i_offset`` /
``grid.j_offset`` by the integrators — is added to the in-memory,
subset-relative positions. A non-zero offset is also stored in the run
metadata as ``i_offset`` / ``j_offset``.
Run metadata (GCM adapter name, integration strategy, etc.) is stored as
a JSON object under the Parquet schema metadata key ``"tracmass.run"``.
Read it back without loading row data::
import json, pyarrow.parquet as pq
meta = json.loads(pq.read_schema("traj.parquet").metadata[b"tracmass.run"])
Parameters
----------
filepath : str
Output ``.parquet`` file path. The file is created on the first
write.
metadata : dict, optional
Arbitrary JSON-serialisable key/value pairs stored in the Parquet
schema metadata under the key ``"tracmass.run"``. Typical fields:
* ``"gcm"`` — GCM adapter name, e.g. ``"glorys"`` or ``"theoretical"``
* ``"integration"`` — function used, e.g. ``"time_dependent"``
* ``"n_windows"`` — number of GCM windows
* ``"no_wflux"`` — whether vertical fluxes were disabled
Attributes
----------
grid_offset : tuple[int, int]
``(i_offset, j_offset)`` added to ``x``/``ia`` and ``y``/``ja`` on
write. Default ``(0, 0)``.
metadata : dict
Run metadata written with the file (a copy of *metadata*); may be
extended until the first write.
Examples
--------
>>> meta = dict(gcm="glorys", integration="time_dependent",
... n_windows=50, no_wflux=True)
>>> with ParticleRecorder("traj.parquet", metadata=meta) as rec:
... result = tracmass.time_dependent(
... state, grid, lazy, fluxes_fn, n_windows=50, recorder=rec,
... no_wflux=True,
... )
"""
_SCHEMA = None # class-level schema cache, populated lazily on first use
[docs]
def __init__(self, filepath: str, metadata: dict | None = None):
self._filepath = filepath
self._writer = None
self.metadata: dict = dict(metadata or {})
self.grid_offset: tuple[int, int] = (0, 0)
def __enter__(self):
return self
def __exit__(self, *_):
self.close()
@classmethod
def _schema(cls):
import pyarrow as pa
if cls._SCHEMA is None:
cls._SCHEMA = pa.schema([
pa.field("event_type", pa.string()),
pa.field("window", pa.int32()),
pa.field("dtm", pa.timestamp("us")),
pa.field("id", pa.int32()),
pa.field("x", pa.float64()),
pa.field("y", pa.float64()),
pa.field("z", pa.float64()),
pa.field("ia", pa.int32()),
pa.field("ja", pa.int32()),
pa.field("ka", pa.int32()),
pa.field("tt", pa.float64()),
pa.field("t0", pa.float64()),
pa.field("reason", pa.int16()),
])
return cls._SCHEMA
def _write(
self,
event_type: str,
idx: np.ndarray,
state: ParticleState,
dtm,
window: int,
) -> None:
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
n = len(idx)
if n == 0:
return
ts = pd.Timestamp(dtm) if dtm is not None else pd.NaT
schema = self._schema()
di, dj = self.grid_offset
table = pa.table(
{
"event_type": [event_type] * n,
"window": np.full(n, window, dtype=np.int32),
"dtm": [ts] * n,
"id": idx.astype(np.int32),
"x": np.asarray(state.x)[idx].astype(np.float64) + di,
"y": np.asarray(state.y)[idx].astype(np.float64) + dj,
"z": np.asarray(state.z)[idx].astype(np.float64),
"ia": np.asarray(state.ia)[idx].astype(np.int32) + di,
"ja": np.asarray(state.ja)[idx].astype(np.int32) + dj,
"ka": np.asarray(state.ka)[idx].astype(np.int32),
"tt": np.asarray(state.tt)[idx].astype(np.float64),
"t0": np.asarray(state.t0)[idx].astype(np.float64),
"reason": np.asarray(state.reason)[idx].astype(np.int16),
},
schema=schema,
)
if self._writer is None:
metadata = dict(self.metadata)
if (di, dj) != (0, 0):
metadata.update(i_offset=int(di), j_offset=int(dj))
if metadata:
import json
schema = schema.with_metadata(
{"tracmass.run": json.dumps(metadata)}
)
self._writer = pq.ParquetWriter(self._filepath, schema)
self._writer.write_table(table)
[docs]
def record(
self,
state: ParticleState,
dtm=None,
*,
window: int = 0,
) -> None:
"""Append active-particle positions as a new row group.
Parameters
----------
state : ParticleState
Current particle state; only active particles are written.
dtm : pd.Timestamp or datetime-like, optional
Timestamp stored in the ``dtm`` column. ``None`` → ``NaT``.
window : int, optional
Window index (0-based). Default ``0``.
"""
mask = np.asarray(state.active)
idx = np.where(mask)[0]
self._write("snapshot", idx, state, dtm, window)
[docs]
def record_rampup(
self,
state: ParticleState,
dtm=None,
*,
window: int = 0,
) -> None:
"""Append active-particle positions saved inside a window.
Like :meth:`record`, but with ``event_type="rampup"``: the extra
saves of the ramp-up period (see ``rampup_windows`` in
:func:`tracmass.time_dependent`). The run metadata's
``save_dt_seconds`` gives the time between consecutive saves.
Parameters
----------
state : ParticleState
Current particle state; only active particles are written.
dtm : pd.Timestamp or datetime-like, optional
Time of the save (inside the window).
window : int, optional
Index of the window the save falls in.
"""
idx = np.where(np.asarray(state.active))[0]
self._write("rampup", idx, state, dtm, window)
[docs]
def record_deactivations(
self,
state: ParticleState,
mask: np.ndarray,
dtm=None,
*,
window: int = 0,
) -> None:
"""Append deactivated particles identified by *mask*.
Parameters
----------
state : ParticleState
Particle state after the integration window.
mask : array-like of bool, shape (N,)
Boolean mask selecting the newly-deactivated particles.
Typically ``active_before & (reason != 0)``.
dtm : pd.Timestamp or datetime-like, optional
Timestamp stored in the ``dtm`` column. ``None`` → ``NaT``.
window : int, optional
Window index (0-based). Default ``0``.
"""
idx = np.where(mask)[0]
self._write("deactivated", idx, state, dtm, window)
[docs]
def close(self) -> None:
"""Finalise and close the Parquet file.
Safe to call even if no events were recorded.
"""
if self._writer is not None:
self._writer.close()
self._writer = None
