Source code for tracmass.trajs
"""Trajectory analysis utilities for TRACMASS-JAX.
Provides functions for computing physical properties from trajectory
DataFrames produced by :class:`~tracmass.particle.ParticleRecorder`.
"""
from __future__ import annotations
import numpy as np
import pandas as pd
from tracmass.trajs.coords import lonlat_to_xy, xy_to_lonlat
__all__ = ["speeds", "xy_to_lonlat", "lonlat_to_xy"]
[docs]
def speeds(
df: pd.DataFrame,
dxdy,
) -> pd.DataFrame:
"""Compute inter-snapshot speeds for all particles in a trajectory DataFrame.
For each consecutive pair of snapshots per particle, computes the physical
displacement in metres and the speed in m/s.
The effective cell width used for displacement conversion is
``sqrt(dxdy_cell)``, which is exact for square cells and the
geometric-mean approximation for non-square cells.
Parameters
----------
df : pd.DataFrame
Trajectory data as produced by
:class:`~tracmass.particle.ParticleRecorder`. Required columns:
``id``, ``x``, ``y``, ``tt``, ``window``. If an ``event_type``
column is present, only ``"snapshot"`` rows are used. If *dxdy*
is a 2-D array, ``ia`` and ``ja`` must also be present.
dxdy : float, array-like of shape (imt, jmt), or GridFields
Cell area in m².\n
- **scalar** — uniform grid; cell width is ``sqrt(dxdy)`` in both
directions. Pass e.g. ``250.0 ** 2`` for 250 m Cartesian cells.
- **2-D array** — spatially varying grid; the cell area at the
*starting* position of each interval is looked up using the
1-based ``ia``/``ja`` columns. Pass ``np.asarray(grid.dxdy)``
or any array-like of the right shape.
- **GridFields** (or any object with a ``.dxdy`` attribute) — the
``.dxdy`` attribute is used as the 2-D array. Pass ``grid``
directly.
Returns
-------
pandas.DataFrame
One row per consecutive snapshot pair per particle, with columns:
``id``
Particle identifier.
``window``
Window index of the *later* snapshot.
``dt``
Time elapsed between snapshots in seconds.
``dx_m``
x-displacement in metres (signed).
``dy_m``
y-displacement in metres (signed).
``speed``
Euclidean speed in m/s.
Examples
--------
Uniform grid (e.g. Fabboni test case, 250 m cells):
>>> import pandas as pd
>>> from tracmass.trajs import speeds
>>> df = pd.read_parquet("traj.parquet")
>>> sp = speeds(df, dxdy=250.0 ** 2)
>>> sp["speed"].max()
Spatially varying grid from a GridFields object:
>>> import tracmass
>>> grid, lazy = tracmass.new_flux(tracmass.gcm.glorys.fluxes)
>>> sp = speeds(df, grid)
"""
# Accept GridFields or anything with a .dxdy attribute
if hasattr(dxdy, "dxdy"):
dxdy = np.asarray(dxdy.dxdy)
scalar = np.ndim(dxdy) == 0
if not scalar:
dxdy = np.asarray(dxdy)
# Filter to snapshot events if the column is present
if "event_type" in df.columns:
df = df[df["event_type"] == "snapshot"]
if df.empty:
return pd.DataFrame(columns=["id", "window", "dt", "dx_m", "dy_m", "speed"])
df = df.sort_values(["id", "window"])
grp = df.groupby("id", sort=False)
dx = grp["x"].diff()
dy = grp["y"].diff()
dt = grp["tt"].diff()
if scalar:
cell_size = np.sqrt(float(dxdy))
else:
ia_s = grp["ia"].shift(1).fillna(1).astype(int).values - 1 # 1-based → 0-based
ja_s = grp["ja"].shift(1).fillna(1).astype(int).values - 1
cell_size = np.sqrt(dxdy[ia_s, ja_s])
dx_m = dx * cell_size
dy_m = dy * cell_size
mask = (dt > 0).values
if not mask.any():
return pd.DataFrame(columns=["id", "window", "dt", "dx_m", "dy_m", "speed"])
return pd.DataFrame({
"id": df["id"].values[mask],
"window": df["window"].values[mask],
"dt": dt.values[mask],
"dx_m": dx_m.values[mask],
"dy_m": dy_m.values[mask],
"speed": (np.sqrt(dx_m.values[mask] ** 2 + dy_m.values[mask] ** 2)
/ dt.values[mask]),
})
