"""Settings-driven, single-call entry point for a trajectory run.
:func:`run` strings together the usual steps — activate a ``tracmass.toml``
environment, build the grid and flux buffer, seed particles, integrate —
taking everything not passed explicitly from the active settings:
================= ==========================================================
Setting Used for
================= ==========================================================
``gcm.name`` GCM adapter (``[env.gcm] name``), when *gcm* is omitted
``dtm_start`` Start date, when *dtm_start* is omitted
``dtm_end`` End date, when *dtm_end* is omitted
``runtime`` ``dtm_end = dtm_start + runtime`` when no end date is set;
a negative value (``"-180D"``) gives a backward run
``integrator`` ``"time_dependent"`` (default), ``"time_dependent_patched"``
or ``"time_dependent_streaming"``, when *integrator* is
omitted
``seed`` Seeding table (:func:`tracmass.seed.from_settings`), when
*state* is omitted
``verbose`` Debug output of each run stage (:mod:`tracmass.verbose`),
when *verbose* is omitted
================= ==========================================================
The integrators themselves resolve ``monitor``, ``recorder``, ``no_wflux`` and
``n_particle_batches`` from settings (see :mod:`tracmass.runsetup`), and
``n_windows`` from the flux buffer's ``dtm_end``.
"""
from __future__ import annotations
import dataclasses
import logging
from typing import Callable
import pandas as pd
from tracmass import config, seed
from tracmass.gcm import GridFields, new_flux
from tracmass.integrate import (
time_dependent,
time_dependent_patched,
time_dependent_streaming,
)
from tracmass.particle import ParticleState
from tracmass.verbose import timed, verbose_output
log = logging.getLogger(__name__)
__all__ = ["run"]
INTEGRATORS: dict[str, Callable] = {
"time_dependent": time_dependent,
"time_dependent_patched": time_dependent_patched,
"time_dependent_streaming": time_dependent_streaming,
}
[docs]
def run(
dtm_start: pd.Timestamp | str | None = None,
dtm_end: pd.Timestamp | str | None = None,
state: ParticleState | Callable[[GridFields], ParticleState] | None = None,
*,
env: str | None = None,
gcm: str | Callable | None = None,
integrator: str | None = None,
verbose: bool | None = None,
**kwargs,
) -> ParticleState:
"""Build the grid and flux buffer, seed, and integrate in one call.
Parameters
----------
dtm_start : str or pd.Timestamp or None, optional
Start date. Default: the ``dtm_start`` setting, else the adapter's
own default date.
dtm_end : str or pd.Timestamp or None, optional
End date; earlier than the start date means a backward run.
Default: the ``dtm_end`` setting, else start + the ``runtime``
setting, else a single window.
state : ParticleState or callable or None, optional
Initial particles, or ``seed(grid) -> ParticleState`` called once the
grid is built (and *env* is active, so it can read settings).
Default: seed from the ``[env.seed]`` settings table via
:func:`tracmass.seed.from_settings`.
env : str or None, optional
``tracmass.toml`` environment to activate first (via
``settings.setenv``; it stays active after the call, which is what
adapters and integrators read during the run). Default: keep the
current environment.
gcm : str or callable or None, optional
Adapter name or ``fluxes`` callable. Default: the ``gcm.name``
setting.
integrator : str or None, optional
Name of the integrator, a key of :data:`INTEGRATORS`. Default: the
``integrator`` setting, else ``"time_dependent"``.
verbose : bool or None, optional
Print timestamped debug messages for each stage (adapter calls,
flux reads and their size, seeding, compilation, each window) to
stderr while the run lasts; see :mod:`tracmass.verbose`. Default:
the ``verbose`` setting, else False.
**kwargs
``dtm`` is accepted as an alias for *dtm_start* (matching the
``dtm`` naming of :class:`~tracmass.gcm.LazyFluxFields`); the rest
are passed on to the integrator (e.g. ``iter_sub``, ``monitor``,
``recorder``, ``patch_buffer``, ``n_windows``).
Returns
-------
ParticleState
Final particle state.
Raises
------
ValueError
If no GCM is given or configured, *integrator* is unknown, *state*
is omitted and there is no usable ``[env.seed]`` table, or a string
*dtm_start* / *dtm_end* is not a date (e.g. ``run("omi")`` where
``run(env="omi")`` was meant).
TypeError
If both *dtm_start* and its alias ``dtm`` are given.
If a particle state or seed callable is passed as *dtm_start* or
*dtm_end* (the pre-2026-09-30 argument order).
Examples
--------
>>> from tracmass import seed
>>> result = tracmass.run(
... "2020-07-01",
... state=lambda grid: seed.ijk(10, 20, 10, 20, 1, 1, grid=grid),
... env="tempvel", integrator="time_dependent_patched",
... )
>>> result = tracmass.run("2007-06-01", env="omi") # seeding from settings too
>>> result = tracmass.run(dtm="2007-06-01", env="omi") # same, via the alias
"""
if "dtm" in kwargs:
if dtm_start is not None:
raise TypeError(
"tracmass.run: got both dtm_start and its alias dtm; pass one"
)
dtm_start = kwargs.pop("dtm")
for name, value in (("dtm_start", dtm_start), ("dtm_end", dtm_end)):
if isinstance(value, ParticleState) or callable(value):
raise TypeError(
f"tracmass.run: got a particle state / seed function as "
f"{name}; the argument order is run(dtm_start, dtm_end, state)"
)
if isinstance(value, str):
try:
pd.Timestamp(value)
except ValueError:
raise ValueError(
f"tracmass.run: cannot parse {name}={value!r} as a date; "
f"if it is a tracmass.toml environment, pass "
f"env={value!r} (env is keyword-only)"
) from None
settings = config.settings
if env is not None:
settings.setenv(env)
with verbose_output(verbose):
log.debug("run: env=%s", settings.current_env)
return _run(dtm_start, dtm_end, state, gcm, integrator, kwargs)
def _run(dtm_start, dtm_end, state, gcm, integrator, kwargs) -> ParticleState:
"""Body of :func:`run` once the environment is active."""
settings = config.settings
if gcm is None:
gcm = config.gcm_name()
if gcm is None:
raise ValueError(
"tracmass.run: no GCM given; pass gcm=... or set [env.gcm] name"
)
if integrator is None:
integrator = settings.get("integrator", None) or "time_dependent"
try:
integrate = INTEGRATORS[integrator]
except KeyError:
raise ValueError(
f"tracmass.run: unknown integrator {integrator!r}; "
f"expected one of {sorted(INTEGRATORS)}"
) from None
if dtm_start is None:
dtm_start = settings.get("dtm_start", None)
if dtm_end is None:
dtm_end = settings.get("dtm_end", None)
runtime = settings.get("runtime", None)
log.debug(
"run: gcm=%s integrator=%s dtm_start=%s dtm_end=%s runtime=%s "
"subset=%s", getattr(gcm, "__name__", gcm), integrator, dtm_start,
dtm_end, runtime, config.gcm_subset(),
)
with timed(log, "run: building grid and flux buffer"):
grid, lazy = _new_flux(gcm, dtm_start, dtm_end, runtime)
log.debug(
"run: grid imt=%d jmt=%d km=%d offset=(%d, %d) timestep=%s "
"reverse=%s dtm=%s n_windows=%s", grid.imt, grid.jmt, grid.km,
grid.i_offset, grid.j_offset, grid.timestep, grid.reverse, lazy.dtm,
lazy.n_windows,
)
with timed(log, "run: seeding"):
if state is None:
state = seed.from_settings(grid)
elif callable(state):
state = state(grid)
log.debug("run: %d particles seeded; calling %s", len(state.x), integrator)
return integrate(state, grid, lazy, **kwargs)
def _new_flux(gcm, dtm_start, dtm_end, runtime):
"""``(grid, lazy)`` for :func:`run`, with *dtm_end* from *runtime*."""
if dtm_end is None and runtime is not None:
runtime = pd.Timedelta(runtime)
if dtm_start is not None:
dtm_end = pd.Timestamp(dtm_start) + runtime
return new_flux(gcm, dtm_start=dtm_start, dtm_end=dtm_end)
# Start date is the adapter's default: build forward to find it,
# then attach dtm_end (or rebuild backward for a negative runtime).
grid, lazy = new_flux(gcm)
dtm_end = lazy.dtm + runtime
if runtime >= pd.Timedelta(0):
return grid, dataclasses.replace(lazy, dtm_end=dtm_end)
return new_flux(gcm, dtm_start=lazy.dtm, dtm_end=dtm_end)
return new_flux(gcm, dtm_start=dtm_start, dtm_end=dtm_end)
