"""Where trajectory files live, and loading them back.
The settings-driven recorder (see :mod:`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.
:func:`output_folder` / :func:`output_stem` build that location from a
settings environment, and :func:`load` finds and reads a file from it — by
file name, by environment and start date, or simply the newest one.
"""
from __future__ import annotations
import contextlib
import importlib
import json
import re
from pathlib import Path
import pandas as pd
__all__ = ["load", "output_folder", "output_stem"]
_DIRECTIONS = ("forward", "reverse")
CLEAN_COLUMNS = ["id", "x", "y", "dtm"]
def _env_settings(env: str | None = None):
"""Settings for *env* (without switching to it), else the active ones."""
from tracmass.config import settings
return settings if env is None else settings.from_env(env)
def _env_name(st) -> str:
return (getattr(st, "current_env", "") or "").lower()
[docs]
def output_stem(st) -> str:
"""File-name stem for *st*: ``output_name``, else ``trm_{env}``.
Parameters
----------
st : dynaconf.LazySettings
Settings of one environment.
Returns
-------
str
"""
env = _env_name(st)
return st.get("output_name", None) or (f"trm_{env}" if env else "trm")
[docs]
def output_folder(st) -> Path:
"""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
"""
base = st.get("base_output_dir", None)
folder = Path(base) if base else Path()
if st.get("use_sub_folder", False):
folder = folder / (st.get("output_name", None) or _env_name(st) or "trm")
return folder
[docs]
def output_filename(stem: str, reverse: bool = False, dtm_start=None) -> str:
"""``{stem}_{forward|reverse}[_{date}].parquet``."""
parts = [stem, _DIRECTIONS[bool(reverse)]]
if dtm_start is not None:
parts.append(str(pd.Timestamp(dtm_start).date()))
return f"{'_'.join(parts)}.parquet"
def _search_folders(st) -> list[Path]:
""":func:`output_folder` first, then the other layouts older files use.
Files written before ``use_sub_folder`` (or with it toggled) sit in
``base_output_dir`` itself or in a sub-folder named after the
environment / ``output_name``.
"""
base = st.get("base_output_dir", None)
root = Path(base) if base else Path()
folders = [output_folder(st), root]
for name in (st.get("output_name", None), _env_name(st)):
if name:
folders.append(root / name)
return list(dict.fromkeys(folders))
def _newest(paths) -> Path | None:
paths = [p for p in paths if p.is_file()]
return max(paths, key=lambda p: p.stat().st_mtime) if paths else None
def _resolve_path(env, dtm, reverse) -> Path:
"""Find the file :func:`load` reads when no file name is given."""
st = _env_settings(env)
if env is None and dtm is None:
# Newest file of any run: base_output_dir and its sub-folders.
base = st.get("base_output_dir", None)
root = Path(base) if base else Path()
found = _newest(
p
for p in [*root.glob("*.parquet"), *root.glob("*/*.parquet")]
if reverse is None or f"_{_DIRECTIONS[bool(reverse)]}" in p.stem
)
if found is None:
raise FileNotFoundError(f"tracmass.load: no .parquet files in {root}")
return found
folders, stem = _search_folders(st), output_stem(st)
directions = _DIRECTIONS if reverse is None else (_DIRECTIONS[bool(reverse)],)
if dtm is not None:
names = [output_filename(stem, d == "reverse", dtm) for d in directions]
# First folder holding a match; a forward/reverse pair there is ambiguous.
existing = next(
(
[folder / n for n in names if (folder / n).is_file()]
for folder in folders
if any((folder / n).is_file() for n in names)
),
[],
)
if len(existing) > 1:
raise ValueError(
f"tracmass.load: both a forward and a reverse run start on "
f"{pd.Timestamp(dtm).date()}; pass reverse=False or True"
)
if not existing:
raise FileNotFoundError(
f"tracmass.load: no {' or '.join(names)} in "
+ ", ".join(str(f) for f in folders)
)
return existing[0]
# env without a date: newest file of that environment.
pattern = re.compile(
rf"{re.escape(stem)}_({'|'.join(directions)})(_\d{{4}}-\d{{2}}-\d{{2}})?\.parquet"
)
found = _newest(
p
for folder in folders
for p in folder.glob("*.parquet")
if pattern.fullmatch(p.name)
)
if found is None:
raise FileNotFoundError(
f"tracmass.load: no {stem}_*.parquet files in "
+ ", ".join(str(f) for f in folders)
)
return found
def _gcm_from_settings(st) -> str | None:
"""``[env.gcm] name`` of *st* (or a bare ``gcm = "<name>"``)."""
gcm = st.get("gcm", None)
if hasattr(gcm, "get"):
return gcm.get("name", None)
return gcm if isinstance(gcm, str) else None
@contextlib.contextmanager
def _source_product(name: str, product: str | None):
"""Temporarily point ``tracmass.seagrid.<name>`` at the run's *product*.
The adapter reads its grid (and so lon/lat) for the product currently
in ``settings["product"]``; a file written for another product must be
converted with that product's grid.
"""
try:
source = importlib.import_module(f"tracmass.seagrid.{name}")
except ImportError:
source = None
st = getattr(source, "settings", None)
if not product or not isinstance(st, dict) or st.get("product") in (None, product):
yield
return
saved = st["product"]
st["product"] = product
try:
yield
finally:
st["product"] = saved
def _add_lonlat(df: pd.DataFrame, gcm: str | None, env: str | None) -> None:
"""Add ``lon`` / ``lat`` columns to *df* in place from its ``x`` / ``y``.
The adapter is *gcm*, else the file's ``gcm`` metadata, else ``[env.gcm]
name`` of *env*, the file's ``env`` metadata or the active environment.
The file's ``gcm_product`` selects the adapter's data product (grid).
"""
meta = df.attrs.get("tracmass.run", {})
name = gcm or meta.get("gcm")
if name is None:
name = _gcm_from_settings(_env_settings(env or meta.get("env")))
if name is None:
raise ValueError(
"tracmass.load: lonlat=True needs the GCM; pass gcm=... (no 'gcm' "
"in the file's metadata or [env.gcm] name in the settings)"
)
module = importlib.import_module(f"tracmass.gcm.{name}")
if not hasattr(module, "xy_to_lonlat"):
raise ValueError(
f"tracmass.load: GCM {name!r} has no xy_to_lonlat; lon/lat are "
f"only available for adapters with a geographic grid"
)
# The product only applies to the adapter that wrote the file.
product = meta.get("gcm_product") if name == meta.get("gcm") else None
with _source_product(name, product):
df["lon"], df["lat"] = module.xy_to_lonlat(
df["x"].to_numpy(), df["y"].to_numpy()
)
def _normalize_filters(filters) -> list[list[tuple]] | None:
"""*filters* of :func:`load` as pyarrow DNF: a list of ANDed tuple lists."""
if filters is None:
return None
if isinstance(filters, dict):
return [[
(col, "in", list(val)) if isinstance(val, (list, tuple, set))
else (col, "==", val)
for col, val in filters.items()
]]
filters = list(filters)
if not filters:
return None
if all(isinstance(f, tuple) for f in filters):
return [filters]
if all(isinstance(f, list) for f in filters):
return [[tuple(t) for t in group] for group in filters]
raise TypeError(
"tracmass.load: filters must be a dict, a list of (column, op, value) "
"tuples, or a list of such lists"
)
[docs]
def load(
filename: str | Path | None = None,
*,
env: str | None = None,
dtm=None,
reverse: bool | None = None,
columns: list[str] | None = None,
filters: dict | list | None = None,
lonlat: bool = False,
gcm: str | None = None,
clean: bool = True,
rampup: bool = False,
) -> pd.DataFrame:
"""Load a trajectory file written by :class:`~tracmass.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 :func:`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.
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
"""
if filename is not None:
if env is not None or dtm is not None:
raise TypeError("tracmass.load: pass either filename or env/dtm")
path = Path(filename)
else:
path = _resolve_path(env, dtm, reverse)
import pyarrow.parquet as pq
schema = pq.read_schema(path)
filters = _normalize_filters(filters)
if clean:
if columns is None:
# Files from other writers may lack some of these; keep those present.
columns = [c for c in CLEAN_COLUMNS if c in schema.names]
if rampup and "event_type" in schema.names:
columns.insert(0, "event_type")
if "event_type" in schema.names:
snapshot = (
("event_type", "in", ["snapshot", "rampup"]) if rampup
else ("event_type", "==", "snapshot")
)
filters = [[*group, snapshot] for group in (filters or [[]])]
read = columns
if lonlat and columns is not None:
read = list(dict.fromkeys([*columns, "x", "y"]))
df = pd.read_parquet(path, columns=read, filters=filters)
if filters is not None:
df = df.reset_index(drop=True)
raw = (schema.metadata or {}).get(b"tracmass.run")
df.attrs["path"] = str(path)
df.attrs["tracmass.run"] = json.loads(raw) if raw else {}
if lonlat:
_add_lonlat(df, gcm, env)
if columns is not None:
df = df[[*columns, "lon", "lat"]]
return df
