"""Resolve settings-driven arguments of the integrators.
Covers ``monitor`` / ``recorder`` (lifecycle-managed here via
:func:`manage_run_io`) and ``no_wflux`` (a plain bool derived from the
``move_vertically`` setting via :func:`resolve_no_wflux`).
The multi-window integration routines
(:func:`~tracmass.integrate.time_dependent`,
:func:`~tracmass.integrate.time_dependent_patched`,
:func:`~tracmass.integrate.time_dependent_streaming`) each accept a
``monitor`` and a ``recorder`` argument. Historically the caller had to
build a :class:`~tracmass.RunMonitor` / :class:`~tracmass.ParticleRecorder`
and drive its context manager by hand::
with (RunMonitor(backend="rich") as mon,
ParticleRecorder("run.parquet") as rec):
tracmass.time_dependent_patched(..., monitor=mon, recorder=rec)
The :func:`manage_run_io` decorator removes that boilerplate. Applied to an
integrator it accepts the same two arguments in any of four forms and owns
the lifecycle of whatever it constructs itself:
==================== ===================================================
value behaviour
==================== ===================================================
an instance used as-is; the *caller* owns its lifecycle
a ``str`` ``RunMonitor(backend=...)`` / ``ParticleRecorder(...)``
built here and entered on an ``ExitStack``
``None`` feature disabled
omitted (``_UNSET``) built from :mod:`tracmass.config` settings and
entered on the ``ExitStack``
==================== ===================================================
"Entered on the ``ExitStack``" means the object's ``__exit__`` runs
automatically when the integrator returns *or raises*, so a ``with`` block
around the call is no longer required.
Settings keys consulted for the omitted (``_UNSET``) case
---------------------------------------------------------
``runmonitor_backend``
Backend for the auto-built :class:`RunMonitor` (``"print"``, ``"log"``,
or ``"rich"``). Default ``"print"``.
``output_name``
Filename stem for the auto-built recorder. Setting it also *enables*
auto-recording. Default stem when auto-recording via ``auto_record``:
``trm_{active-env}``.
``auto_record``
Truthy value enables auto-recording even when ``output_name`` is unset.
``base_output_dir``
Directory the auto-built recorder writes into. Default: current
directory.
``use_sub_folder``
When truthy, the auto-built recorder writes into a sub-directory of
``base_output_dir`` named after ``output_name``, else the active
environment. Default ``False``. Missing directories are created.
The auto-built filename is
``{stem}_{forward|reverse}_{dtm_start.date()}.parquet`` — the direction token
comes from ``grid.reverse`` and the date from ``fluxes.dtm``.
When neither ``output_name`` nor ``auto_record`` is set, an omitted
``recorder`` resolves to ``None`` (no trajectory file) — the historical
default.
Vertical advection
------------------
``move_vertically``
When an integrator's ``no_wflux`` argument is not passed explicitly,
:func:`resolve_no_wflux` sets ``no_wflux = not move_vertically`` from
this setting (default ``True`` — particles do move vertically, matching
the historical ``no_wflux=False`` default). An explicit ``no_wflux=``
always wins. The ECCO2 adapter reads the same setting for its
``include_wflux`` argument (see :func:`tracmass.config.move_vertically`).
Particle batching
-----------------
``n_particle_batches``
When an integrator's ``n_particle_batches`` argument is not passed
explicitly, :func:`resolve_n_particle_batches` takes it from this
setting (default ``1`` — no batching). An explicit value always wins;
both are clamped to ``>= 1``. The full-grid integrators
(:func:`~tracmass.integrate.steady_state`,
:func:`~tracmass.integrate.time_dependent`) act on it with a
particle-batched ``lax.while_loop``; the streaming integrators accept it
for API symmetry only.
Ramp-up saves
-------------
``rampup_windows`` / ``rampup_substeps``
When an integrator's ``rampup_windows`` / ``rampup_substeps`` arguments
are not passed, :func:`resolve_rampup` takes them from these settings
(defaults ``0`` — no ramp-up — and ``5``). During the first
``rampup_windows`` windows positions are also saved
``rampup_substeps - 1`` times inside each window, as
``event_type="rampup"`` rows; :func:`rampup_metadata` describes the
resulting save times in the trajectory file.
"""
from __future__ import annotations
import contextlib
import functools
import re
from pathlib import Path
from typing import Callable
from tracmass.monitoring import RunMonitor
from tracmass.particle import ParticleRecorder
from tracmass.verbose import verbose_output
__all__ = [
"manage_run_io",
"resolve_monitor",
"resolve_n_particle_batches",
"resolve_no_wflux",
"resolve_rampup",
"resolve_recorder",
"resolve_seed_freq",
"resolve_timax",
"rampup_metadata",
"reseed",
"run_metadata",
"seed_metadata",
"seed_windows",
"traj_max_seconds",
"window_ends",
]
TIMAX_NONE = 1e20 # timax_seconds when there is no lifetime limit
# Sentinel: the argument was not supplied at all (distinct from explicit None).
_UNSET: object = object()
[docs]
def resolve_no_wflux(no_wflux) -> bool:
"""Normalise an integrator ``no_wflux`` argument to a plain bool.
Parameters
----------
no_wflux : bool or _UNSET
An explicit bool is returned unchanged (the caller always wins).
``_UNSET`` (the argument was omitted) is resolved from settings as
``not move_vertically()`` — see
:func:`tracmass.config.move_vertically`.
Returns
-------
bool
"""
if no_wflux is not _UNSET:
return bool(no_wflux)
try:
from tracmass.config import move_vertically
return not move_vertically()
except Exception: # pragma: no cover - config is optional
return False
[docs]
def resolve_n_particle_batches(n_particle_batches) -> int:
"""Normalise an integrator ``n_particle_batches`` argument to a plain int.
Parameters
----------
n_particle_batches : int or _UNSET
An explicit value is returned as ``max(1, int(...))`` (the caller
always wins). ``_UNSET`` (the argument was omitted) is resolved from
the ``n_particle_batches`` setting — see
:func:`tracmass.config.n_particle_batches` (default ``1``).
Returns
-------
int
Batching factor, clamped to ``>= 1``.
"""
if n_particle_batches is not _UNSET:
try:
return max(1, int(n_particle_batches))
except (TypeError, ValueError):
return 1
try:
from tracmass.config import n_particle_batches as _setting
return _setting()
except Exception: # pragma: no cover - config is optional
return 1
[docs]
def resolve_rampup(rampup_windows, rampup_substeps) -> tuple[int, int]:
"""Normalise the integrators' ramp-up arguments to ``(windows, substeps)``.
Parameters
----------
rampup_windows : int or _UNSET
Number of leading windows saved more often. ``_UNSET`` → the
``rampup_windows`` setting, default ``0`` (no ramp-up).
rampup_substeps : int or _UNSET
Saves per window during ramp-up (``5`` → every ``timestep / 5``).
``_UNSET`` → the ``rampup_substeps`` setting, default ``5``.
Returns
-------
tuple[int, int]
``(rampup_windows, rampup_substeps)``; ``(0, 1)`` when ramp-up is
off (either is ``0`` / ``1``).
Raises
------
ValueError
If either is negative / below 1, or not an integer.
"""
st = _settings()
if rampup_windows is _UNSET:
rampup_windows = st.get("rampup_windows", 0) if st is not None else 0
if rampup_substeps is _UNSET:
rampup_substeps = st.get("rampup_substeps", 5) if st is not None else 5
windows = int(rampup_windows or 0)
substeps = int(5 if rampup_substeps is None else rampup_substeps)
if windows < 0 or substeps < 1:
raise ValueError(
f"rampup_windows must be >= 0 and rampup_substeps >= 1, got "
f"{rampup_windows!r} / {rampup_substeps!r}"
)
if windows == 0 or substeps == 1:
return 0, 1
return windows, substeps
[docs]
def window_ends(win: int, rampup_windows: int, rampup_substeps: int) -> list[float]:
"""Normalised times at which window *win* stops to save positions.
``[win + 1]`` normally; ``[win + 1/k, win + 2/k, ..., win + 1]`` with
``k = rampup_substeps`` during the first *rampup_windows* windows. All
but the last are ramp-up saves; the last is the next window's snapshot.
"""
if win < rampup_windows:
k = rampup_substeps
return [win + m / k for m in range(1, k)] + [float(win + 1)]
return [float(win + 1)]
[docs]
def resolve_n_windows(n_windows, fluxes) -> int:
"""Normalise an integrator ``n_windows`` argument to a plain int.
Parameters
----------
n_windows : int or None
An explicit value is returned unchanged (the caller always wins).
``None`` is resolved from ``fluxes.n_windows`` — the number of
windows from ``fluxes.dtm`` to the *dtm_end* given to
:func:`tracmass.gcm.new_flux` — falling back to ``1`` when no
*dtm_end* was given (or *fluxes* is not a ``LazyFluxFields``).
fluxes : LazyFluxFields or FluxFields
The flux buffer passed to the integrator.
Returns
-------
int
"""
if n_windows is not None:
return int(n_windows)
derived = getattr(fluxes, "n_windows", None)
return 1 if derived is None else derived
[docs]
def resolve_seed_freq(seed_freq) -> str | None:
"""Normalise an integrator ``seed_freq`` argument.
Parameters
----------
seed_freq : str or None or _UNSET
A pandas frequency (``"30D"``, ``"MS"``, ...) or ``None`` (seed
once) is returned unchanged; the caller always wins. ``_UNSET``
(the argument was omitted) is resolved from ``seed_freq`` in the
``[env.seed]`` settings table, default ``None``.
Returns
-------
str or None
The frequency, or ``None`` when particles are seeded only at the
start (also for an empty string or ``false``).
"""
if seed_freq is _UNSET:
st = _settings()
cfg = st.get("seed", None) if st is not None else None
seed_freq = None
if cfg:
kw = {str(k).lower(): v for k, v in dict(cfg).items()}
seed_freq = kw.get("seed_freq")
return seed_freq or None
[docs]
def seed_windows(
seed_freq: str | None,
dtm_start,
timestep,
n_windows: int,
reverse: bool = False,
) -> list[int]:
"""Windows at whose start a repeat seed is released.
The seed times are *dtm_start* stepped by the pandas offset
*seed_freq* — forward in time, or backward for a backward run — and each
is placed on the window that starts at it. A seed time inside a window
moves to the start of the next window. The initial seed (window 0) is
not included.
Parameters
----------
seed_freq : str or None
Pandas frequency string, e.g. ``"30D"``, ``"7D"``, ``"MS"`` (month
starts) or ``"QS"``. ``None`` → no repeat seeds.
dtm_start : str or pd.Timestamp
Start time of window 0.
timestep : pd.Timedelta or datetime.timedelta
Window length.
n_windows : int
Number of windows in the run; seeds at or after its end are dropped.
reverse : bool, optional
Backward run (windows step back in time). Default ``False``.
Returns
-------
list of int
Sorted, unique window indices, all in ``1 .. n_windows - 1``.
Raises
------
ValueError
If *seed_freq* is not a valid pandas frequency.
Examples
--------
>>> seed_windows("10D", "2010-06-01", pd.Timedelta("1D"), 30)
[10, 20]
>>> seed_windows("MS", "2010-06-15", pd.Timedelta("1D"), 40)
[16]
"""
import warnings
import numpy as np
import pandas as pd
if not seed_freq:
return []
try:
offset = pd.tseries.frequencies.to_offset(seed_freq)
except ValueError as err:
raise ValueError(
f"seed_freq {seed_freq!r} is not a pandas frequency "
f"(e.g. '30D', '7D', 'MS')"
) from err
start = pd.Timestamp(dtm_start)
step = pd.Timedelta(timestep)
end = start + step * n_windows if not reverse else start - step * n_windows
windows: list[int] = []
snapped = False
t = start
while True:
t_next = t - offset if reverse else t + offset
if t_next == t: # a zero offset would loop forever
raise ValueError(f"seed_freq {seed_freq!r} does not advance in time")
t = t_next
if (t >= end) if not reverse else (t <= end):
break
frac = abs(t - start) / step
win = int(np.ceil(frac - 1e-9))
snapped |= win != frac
if 0 < win < n_windows and (not windows or win != windows[-1]):
windows.append(win)
if snapped:
warnings.warn(
f"seed_freq {seed_freq!r} does not fall on window boundaries "
f"(timestep {step}); seeds are released at the start of the next "
f"window",
stacklevel=2,
)
return windows
[docs]
def reseed(state, seed, win: int, tseas: float):
"""Append a copy of *seed* released at the start of window *win*.
The copy's clock is moved forward by *win* windows (``tt`` and ``t0`` by
``win * tseas`` seconds, ``ts`` by *win*) and its iteration counter is
reset, so it starts as a fresh set of particles at that time. Active
flags and reasons come from *seed* (particles seeded inactive stay so).
Parameters
----------
state : ParticleState
Current state of the run.
seed : ParticleState
The run's initial state, as passed to the integrator.
win : int
Window index of the release.
tseas : float
Window length in seconds.
Returns
-------
ParticleState
``ParticleState.concat(state, shifted_seed)``: existing particles
keep their indices (and output ``id``), the new ones follow.
"""
import jax.numpy as jnp
import numpy as np
from tracmass.particle import ParticleState
xp = np if isinstance(state.x, np.ndarray) else jnp
shifted = seed._replace(
tt=seed.tt + xp.asarray(win * tseas, dtype=seed.tt.dtype),
t0=seed.t0 + xp.asarray(win * tseas, dtype=seed.t0.dtype),
ts=seed.ts + xp.asarray(win, dtype=seed.ts.dtype),
niter=xp.zeros_like(seed.niter),
)
if xp is np:
shifted = type(seed)(*(np.asarray(a) for a in shifted))
return ParticleState.concat(state, shifted)
_CALENDAR = re.compile(r"^\s*(\d*)\s*(M|Q|Y)\s*$")
_CALENDAR_MONTHS = {"M": 1, "Q": 3, "Y": 12}
[docs]
def traj_max_seconds(traj_max_time, dtm_start=None, reverse: bool = False) -> float:
"""Convert a trajectory lifetime to seconds.
Parameters
----------
traj_max_time : str or pd.Timedelta or float or None
The longest a particle may live, counted from its release. A
``pd.Timedelta`` string (``"30D"``, ``"12h"``, ``"2W"``), or calendar
months / quarters / years as ``"M"``, ``"Q"``, ``"Y"`` with an
optional count (``"1M"``, ``"6M"``, ``"2Y"``). A number is taken as
seconds; ``None`` means no limit.
dtm_start : str or pd.Timestamp, optional
The run's start time. Calendar units have no fixed length, so they
are measured from here (backward for a backward run): ``"1M"`` from
2010-02-01 is 28 days. Needed only for calendar units.
reverse : bool, optional
Backward run.
Returns
-------
float
Lifetime in seconds (``1e20`` for no limit).
Raises
------
ValueError
If *traj_max_time* cannot be parsed, is not positive, or is in
calendar units without *dtm_start*.
"""
import pandas as pd
if traj_max_time is None or traj_max_time is False or traj_max_time == "":
return TIMAX_NONE
if isinstance(traj_max_time, (int, float)) and not isinstance(
traj_max_time, bool
):
seconds = float(traj_max_time)
else:
match = _CALENDAR.match(str(traj_max_time))
if match:
if dtm_start is None:
raise ValueError(
f"traj_max_time {traj_max_time!r} is in calendar units and "
"needs the run's start time"
)
n = int(match[1] or 1) * _CALENDAR_MONTHS[match[2]]
t0 = pd.Timestamp(dtm_start)
offset = pd.DateOffset(months=n)
seconds = abs(((t0 - offset) if reverse else (t0 + offset)) - t0)
seconds = seconds.total_seconds()
else:
try:
seconds = pd.Timedelta(traj_max_time).total_seconds()
except ValueError as err:
raise ValueError(
f"traj_max_time {traj_max_time!r} is neither a pandas "
"Timedelta (e.g. '30D', '12h') nor calendar units "
"('1M', '1Q', '1Y')"
) from err
if not seconds > 0:
raise ValueError(f"traj_max_time {traj_max_time!r} must be positive")
return seconds
[docs]
def resolve_timax(
timax_seconds, traj_max_time, dtm_start=None, reverse: bool = False
) -> tuple[float, dict]:
"""Resolve the trajectory lifetime of an integrator call.
Parameters
----------
timax_seconds : float or _UNSET
Lifetime in seconds. An explicit value wins.
traj_max_time : str or None or _UNSET
Lifetime as a pandas-style string, see :func:`traj_max_seconds`.
``_UNSET`` is resolved from the ``traj_max_time`` setting, default
no limit.
dtm_start : str or pd.Timestamp, optional
The run's start time, for calendar units.
reverse : bool, optional
Backward run.
Returns
-------
tuple[float, dict]
``timax_seconds`` and run-metadata keys (``traj_max_time`` and
``timax_seconds``; empty without a limit).
Raises
------
ValueError
If both arguments are given explicitly.
"""
if timax_seconds is not _UNSET:
if traj_max_time is not _UNSET and traj_max_time is not None:
raise ValueError("pass either timax_seconds or traj_max_time, not both")
seconds, label = float(timax_seconds), None
else:
if traj_max_time is _UNSET:
st = _settings()
traj_max_time = st.get("traj_max_time", None) if st is not None else None
seconds = traj_max_seconds(traj_max_time, dtm_start, reverse)
label = traj_max_time
if seconds >= TIMAX_NONE:
return TIMAX_NONE, {}
meta = dict(timax_seconds=seconds)
if label is not None:
meta["traj_max_time"] = str(label)
return seconds, meta
def _settings():
"""Return the Dynaconf settings object, or ``None`` if it cannot load."""
try:
from tracmass.config import settings
return settings
except Exception: # pragma: no cover - config is optional
return None
[docs]
def resolve_monitor(monitor, stack: contextlib.ExitStack) -> RunMonitor | None:
"""Normalise *monitor* to a :class:`RunMonitor` instance or ``None``.
Parameters
----------
monitor : RunMonitor or str or None or _UNSET
See the module docstring for the four accepted forms.
stack : contextlib.ExitStack
A monitor constructed here is registered on this stack so its
``__exit__`` runs when the integrator returns or raises. A monitor
passed in by the caller is returned untouched.
Returns
-------
RunMonitor or None
"""
if isinstance(monitor, RunMonitor) or monitor is None:
return monitor
if isinstance(monitor, str):
backend = monitor
else: # _UNSET -> consult settings
st = _settings()
backend = (st.get("runmonitor_backend", "print") if st else "print") or "print"
return stack.enter_context(RunMonitor(backend=backend))
[docs]
def resolve_recorder(
recorder,
stack: contextlib.ExitStack,
*,
dtm_start=None,
reverse: bool = False,
) -> ParticleRecorder | None:
"""Normalise *recorder* to a :class:`ParticleRecorder` instance or ``None``.
Parameters
----------
recorder : ParticleRecorder or str or os.PathLike or None or _UNSET
See the module docstring for the four accepted forms.
stack : contextlib.ExitStack
A recorder constructed here is registered on this stack so its
Parquet writer is closed when the integrator returns or raises. A
recorder passed in by the caller is returned untouched.
dtm_start : pandas.Timestamp or datetime-like, optional
Start timestamp of the run, used only for the auto-built filename.
reverse : bool, optional
Direction of the run (``grid.reverse``). Selects the ``forward`` /
``reverse`` token in the auto-built filename. Default ``False``.
Returns
-------
ParticleRecorder or None
"""
if isinstance(recorder, ParticleRecorder) or recorder is None:
return recorder
if recorder is _UNSET:
path = _auto_recorder_path(dtm_start, reverse)
if path is None:
return None
path.parent.mkdir(parents=True, exist_ok=True)
else: # str / PathLike
path = Path(recorder)
return stack.enter_context(ParticleRecorder(str(path)))
def _auto_recorder_path(dtm_start, reverse: bool = False) -> Path | None:
"""Build the settings-driven Parquet path, or ``None`` if not enabled.
Auto-recording is enabled when ``output_name`` is set or ``auto_record``
is truthy in the active settings environment. The filename is
``{stem}_{forward|reverse}_{dtm_start.date()}.parquet``
(the date is dropped when *dtm_start* is ``None``), where *stem* is
``output_name`` or ``trm_{active-env}``, written into ``base_output_dir``
(current directory when unset) — or, when ``use_sub_folder`` is truthy,
into its sub-directory ``output_name`` (else ``{active-env}``).
"""
from tracmass.output import output_filename, output_folder, output_stem
st = _settings()
if st is None:
return None
if not (st.get("output_name", None) or st.get("auto_record", False)):
return None
return output_folder(st) / output_filename(output_stem(st), reverse, dtm_start)
[docs]
def manage_run_io(fn: Callable) -> Callable:
"""Decorate an integrator to own its *monitor* / *recorder* lifecycle.
The wrapped function is always called with ``monitor`` and ``recorder``
already normalised to an instance or ``None`` (via :func:`resolve_monitor`
/ :func:`resolve_recorder`), with verbose debug output on if the
``verbose`` setting is true (:mod:`tracmass.verbose`), wrapped in a :class:`contextlib.ExitStack`
that closes any object constructed on the caller's behalf when the
integrator returns or raises. The function body no longer interprets
sentinels or constructs these objects itself.
Parameters
----------
fn : callable
An integrator called as ``fn(state, grid, fluxes, ...)`` — ``grid``
supplies the run direction (``grid.reverse``) and ``fluxes`` the run
start timestamp (``fluxes.dtm``) for the auto-built recorder filename
— and which accepts keyword-only ``monitor`` and ``recorder``
parameters. The recorder's ``grid_offset`` is set from
``grid.i_offset`` / ``grid.j_offset`` so it writes full-grid
positions, and :func:`run_metadata` is added to its metadata (keys
already there are kept).
Returns
-------
callable
Wrapper carrying *fn*'s name, docstring and signature via
:func:`functools.wraps`.
"""
@functools.wraps(fn)
def wrapper(*args, monitor=_UNSET, recorder=_UNSET, **kwargs):
grid = args[1] if len(args) >= 2 else kwargs.get("grid")
fluxes = args[2] if len(args) >= 3 else kwargs.get("fluxes")
dtm_start = getattr(fluxes, "dtm", None)
reverse = bool(getattr(grid, "reverse", False))
with contextlib.ExitStack() as stack:
stack.enter_context(verbose_output())
monitor = resolve_monitor(monitor, stack)
recorder = resolve_recorder(
recorder, stack, dtm_start=dtm_start, reverse=reverse
)
if recorder is not None:
# Saved positions are full-grid indices even on a subset grid.
recorder.grid_offset = (
int(getattr(grid, "i_offset", 0)),
int(getattr(grid, "j_offset", 0)),
)
# Describe the run so tracmass.load can find the adapter for
# lon/lat; values the caller already put there win.
for key, value in run_metadata(fn, grid, fluxes).items():
recorder.metadata.setdefault(key, value)
return fn(*args, monitor=monitor, recorder=recorder, **kwargs)
return wrapper
