tracmass.output, particle, monitoring#

Where trajectory files live, and loading them back.

The settings-driven recorder (see tracmass.runsetup) writes to

{base_output_dir}[/{output_name or env}]/{stem}_{forward|reverse}_{date}.parquet

where stem is output_name or trm_{env}, the sub-folder is present when use_sub_folder is truthy, and date is the run’s start date. output_folder() / output_stem() build that location from a settings environment, and load() finds and reads a file from it — by file name, by environment and start date, or simply the newest one.

tracmass.output.output_stem(st)[source]#

File-name stem for st: output_name, else trm_{env}.

Parameters:

st (dynaconf.LazySettings) – Settings of one environment.

Returns:

str

Return type:

str

tracmass.output.output_folder(st)[source]#

Directory trajectory files of st are written to.

base_output_dir (current directory when unset), plus a sub-folder named output_name (else the environment) when use_sub_folder is truthy.

Parameters:

st (dynaconf.LazySettings) – Settings of one environment.

Returns:

Path

Return type:

Path

tracmass.output.output_filename(stem, reverse=False, dtm_start=None)[source]#

{stem}_{forward|reverse}[_{date}].parquet.

Parameters:
Return type:

str

tracmass.output.load(filename=None, *, env=None, dtm=None, reverse=None, columns=None, filters=None, lonlat=False, gcm=None, clean=True, rampup=False)[source]#

Load a trajectory file written by ParticleRecorder.

The file is chosen by, in order of precedence:

  1. filename, if given;

  2. env and/or dtm: the file the settings-driven recorder writes for that environment (default: the active one) and run start date — {output_folder}/{stem}_{forward|reverse}_{dtm}.parquet (also looked for in base_output_dir itself and its {env} / {output_name} sub-folders, where files written with a different use_sub_folder setting end up). Without dtm, the newest file of that environment;

  3. nothing given: the newest .parquet file in base_output_dir or one of its sub-folders, i.e. the most recently written run.

Parameters:

  • filename (str or Path or None, optional) – Path of the file to read.

  • env (str or None, optional) – tracmass.toml environment whose base_output_dir / use_sub_folder / output_name locate the file. Read with settings.from_env; the active environment is not changed.

  • dtm (str or pd.Timestamp or None, optional) – Run start date (the date in the file name).

  • reverse (bool or None, optional) – Run direction. Default None: either, as long as only one matches.

  • columns (list of str or None, optional) – Columns to read. Default: all.

  • filters (dict or list or None, optional) – Rows to read, applied while reading (row groups that cannot match are skipped). A dict {"window": 0} means column == value for each item (a list/tuple/set value means column in value); otherwise pyarrow filters: a list of (column, op, value) tuples (ANDed), or a list of such lists (ORed). Filter columns need not be in columns. Combined (AND) with the clean row filter. Default: all rows.

  • lonlat (bool, optional) – Add lon / lat columns (degrees), converted from x / y with the GCM adapter’s xy_to_lonlat (glorys, ecco2). x / y are read for this even if not in columns. Default False.

  • gcm (str or None, optional) – Adapter for lonlat. Default: gcm in the file’s metadata, else [env.gcm] name of env, of the env in the file’s metadata, or of the active environment. The file’s gcm_product (e.g. "glorys12v1") selects which of the adapter’s grids is used, whatever product is currently configured.

  • clean (bool, optional) – Return only the snapshot rows (event_type == "snapshot"; no deactivation events) and only the columns id, x, y, dtm (plus lon, lat with lonlat). An explicit columns takes precedence over the column part. Default True; clean=False returns every row and column.

  • rampup (bool, optional) – With clean, also keep the event_type == "rampup" rows (the extra saves inside the first windows of a run with rampup_windows), and add the event_type column (unless columns is given) to tell them from the window snapshots. The time between consecutive saves is df.attrs["tracmass.run"] ["save_dt_seconds"]. Default False.

Returns:

pd.DataFrame – The trajectory rows. df.attrs["path"] is the file read and df.attrs["tracmass.run"] its run metadata ({} if none; see tracmass.runsetup.run_metadata() for the keys).

Raises:

  • FileNotFoundError – If no matching file exists.

  • TypeError – If filename is given together with env / dtm.

  • ValueError – If dtm matches both a forward and a reverse run and reverse is not given, or lonlat is set and no GCM with xy_to_lonlat can be found.

Return type:

DataFrame

Examples

>>> df = tracmass.load()                          # newest run, id/x/y/dtm
>>> df = tracmass.load(clean=False)               # every row and column
>>> df = tracmass.load(env="omi_falkor")          # newest omi_falkor run
>>> df = tracmass.load(env="omi_falkor", dtm="2007-06-01")
>>> df = tracmass.load("traj.parquet", columns=["id", "dtm", "x", "y"])
>>> df = tracmass.load(filters={"window": 0})     # first window only
>>> df = tracmass.load(filters=[("window", "<", 10), ("reason", "==", 2)])
>>> df = tracmass.load(env="omi_falkor", lonlat=True)   # + lon, lat columns
>>> df = tracmass.load(lonlat=True)  # id, x, y, dtm, lon, lat
>>> df = tracmass.load(rampup=True)  # + ramp-up saves, event_type column
class tracmass.particle.ParticleState(x, y, z, ia, ja, ka, tt, ts, t0, subvol, niter, active, reason)[source]#

Bases: NamedTuple

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

Parameters:
x: Array#

Alias for field number 0

y: Array#

Alias for field number 1

z: Array#

Alias for field number 2

ia: Array#

Alias for field number 3

ja: Array#

Alias for field number 4

ka: Array#

Alias for field number 5

tt: Array#

Alias for field number 6

ts: Array#

Alias for field number 7

t0: Array#

Alias for field number 8

subvol: Array#

Alias for field number 9

niter: Array#

Alias for field number 10

active: Array#

Alias for field number 11

reason: Array#

Alias for field number 12

classmethod concat(*states)[source]#

Join particle states end to end.

Particles keep their order, so the particles of the first state keep their indices (the id written by 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.

Return type:

ParticleState

Examples

>>> both = ParticleState.concat(state, new_particles)
>>> len(both.x) == len(state.x) + len(new_particles.x)
True
tracmass.particle.write_trajectories_parquet(filepath, state, nff=1)[source]#

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.

Return type:

str

class tracmass.particle.ParticleRecorder(filepath, metadata=None)[source]#

Bases: object

Streams trajectory snapshots and deactivation events to a Parquet file.

Each call to record() appends active-particle positions as a new row group with event_type="snapshot" (at a window boundary); each call to record_rampup() the same with event_type="rampup" (extra saves inside the first windows of a run). Each call to record_deactivations() appends newly-stopped particles with event_type="deactivated". Call 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: 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

grid_offset#

(i_offset, j_offset) added to x/ia and y/ja on write. Default (0, 0).

Type:

tuple[int, int]

metadata#

Run metadata written with the file (a copy of metadata); may be extended until the first write.

Type:

dict

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,
...     )
record(state, dtm=None, *, window=0)[source]#

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.

Return type:

None

record_rampup(state, dtm=None, *, window=0)[source]#

Append active-particle positions saved inside a window.

Like record(), but with event_type="rampup": the extra saves of the ramp-up period (see rampup_windows in 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.

Return type:

None

record_deactivations(state, mask, dtm=None, *, window=0)[source]#

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.

Return type:

None

close()[source]#

Finalise and close the Parquet file.

Safe to call even if no events were recorded.

Return type:

None

Run monitoring for TRACMASS-JAX integration loops.

Provides RunMonitor, a context-manager that receives structured events from time_dependent() and renders them via one of three backends:

  • "print" — simple print() output, reproducing the original behaviour.

  • "log" — structured logging records (one per window + summary).

  • "rich" — live terminal dashboard powered by rich.

tracmass.monitoring.sec_str(s)[source]#

Format a duration in seconds as a compact human-readable string.

Parameters:

s (float) – Duration in seconds.

Returns:

str – E.g. "42ms", "3.7s", "2m05s".

Return type:

str

tracmass.monitoring.dt_str(dt)[source]#

Format a pandas.Timedelta as a compact string.

Parameters:

dt (pandas.Timedelta)

Returns:

str

Return type:

str

class tracmass.monitoring.RunMonitor(backend='print', logger=None, max_table_rows=10)[source]#

Bases: object

Monitor for time_dependent() runs.

Receives structured events from the integration loop and renders them via the chosen backend. Use as a context manager so the "rich" backend can properly start and stop its live display.

Parameters:

  • backend ({"print", "log", "rich"}) – Output backend. "print" reproduces the original stdout output. "log" emits logging records. "rich" shows a live terminal dashboard.

  • logger (logging.Logger, optional) – Logger used by the "log" backend. Defaults to the "tracmass" logger.

  • max_table_rows (int, optional) – Maximum number of window rows shown in the "rich" table at any one time. Older rows scroll out of view. Default 10.

Examples

>>> with tracmass.RunMonitor(backend="rich") as mon:
...     result = tracmass.time_dependent(state, grid, fluxes,
...                                      n_windows=50, monitor=mon)

>>> # Plain print (default when monitor is not supplied):
>>> result = tracmass.time_dependent(state, grid, fluxes, n_windows=50)
on_run_start(n_windows, n_particles)[source]#

Called once before the integration loop begins.

Parameters:

  • n_windows (int) – Total number of GCM windows to integrate.

  • n_particles (int) – Initial number of particles.

Return type:

None

on_window_start(win, n_active, dtm)[source]#

Called at the start of each window, before I/O and advection.

Parameters:

  • win (int) – Zero-based window index.

  • n_active (int) – Number of active particles entering this window.

  • dtm (pandas.Timestamp) – Timestamp of the current GCM time window.

Return type:

None

on_window_end(win, dtm, n_active_before, n_active_after, n_stopped, reason_tally, n_sub, limit_hit, t_io, t_jax, t_py)[source]#

Called at the end of each window after advection and reactivation.

Parameters:

  • win (int) – Zero-based window index.

  • dtm (pandas.Timestamp) – Timestamp of the window just completed.

  • n_active_before (int) – Active count at window entry.

  • n_active_after (int) – Active count after advection + reactivation.

  • n_stopped (int) – Particles permanently deactivated this window.

  • reason_tally (dict[str, int]) – Counts per deactivation reason name.

  • n_sub (int) – Number of sub-cycles executed.

  • limit_hit (bool) – Whether the sub-cycle limit was reached.

  • t_io (float) – Wall time spent in I/O (Dask/Zarr reads), seconds.

  • t_jax (float) – Wall time spent in JAX advection (device_put + while_loop), seconds.

  • t_py (float) – Wall time spent in Python overhead (reactivation, GC, flux swap), seconds.

Return type:

None

on_run_end()[source]#

Called once after the integration loop finishes.

Return type:

None

Verbose debug output: timestamped progress messages for each run stage.

tracmass logs what it is doing (adapter calls, flux reads with their size, seeding, kernel compilation, each window’s phases) to the standard logging logger "tracmass" at DEBUG level. Nothing is shown unless verbose output is turned on, in any of these ways:

  • tracmass.run(..., verbose=True) or verbose = true in the active tracmass.toml environment (also honoured by the integrators);

  • set_verbose() for the rest of the session;

  • standard logging config, e.g. logging.getLogger("tracmass").setLevel("DEBUG") plus a handler.

Messages go to stderr, one line each, with the wall-clock time, so a run that seems stuck shows which step it is in and how long the previous steps took.

tracmass.verbose.set_verbose(verbose=True)[source]#

Turn verbose debug output on or off for the session.

Parameters:

verbose (bool, optional) – True attaches a stderr handler to the "tracmass" logger and sets it to DEBUG; False removes the handler again and restores the previous level. Default True.

Return type:

None

Examples

>>> tracmass.set_verbose()
>>> result = tracmass.run(env="gog")
tracmass.verbose.is_verbose()[source]#

True while set_verbose() output is on.

Return type:

bool

tracmass.verbose.resolve_verbose(verbose=None)[source]#

verbose if given, else the verbose setting (default False).

Parameters:

verbose (bool | None)

Return type:

bool

tracmass.verbose.verbose_output(verbose=None)[source]#

Turn verbose output on for a with block, if verbose resolves True.

Output that was already on (from set_verbose() or an outer block) stays on afterwards; output turned on here is turned off on exit.

Parameters:

verbose (bool | None)

tracmass.verbose.timed(log, msg, *args)[source]#

Log msg ... at entry and msg: done in N s at exit (DEBUG).

Parameters:
tracmass.verbose.gb(nbytes)[source]#

nbytes formatted as "1.23 GB".

Parameters:

nbytes (float)

Return type:

str