Analysing trajectories

Fortran TRACMASS ships its own post-processing (summaries, stream functions, divergences) and the pyTRAJ toolbox. TRACMASS-JAX writes trajectories to Parquet, so the usual Python tools work directly. This page collects common recipes.

Load a run

import tracmass

df = tracmass.load(env="omi", lonlat=True)   # id, x, y, dtm, lon, lat
meta = df.attrs["tracmass.run"]
print(meta["gcm"], meta["direction"], meta["dtm_start"], meta["dtm_end"])

For large files, read only what you need:

df = tracmass.load(env="omi", columns=["id", "dtm", "x", "y"],
                   filters=[("window", "<", 30)])

Plot trajectories

import matplotlib.pyplot as plt

fig, ax = plt.subplots()
for _, traj in df.groupby("id"):
    ax.plot(traj.lon, traj.lat, lw=0.3, alpha=0.5)
start = df[df.dtm == df.dtm.min()]
ax.scatter(start.lon, start.lat, s=2, c="k", label="release")
ax.set_xlabel("longitude"); ax.set_ylabel("latitude"); ax.legend()

For millions of particles, a density map is more useful than lines:

import numpy as np

H, xe, ye = np.histogram2d(df.lon, df.lat, bins=(360, 180))
plt.pcolormesh(xe, ye, np.log10(H.T + 1))

Why particles stopped

full = tracmass.load(env="omi", clean=False)
stops = full[full.event_type == "deactivated"]
print(stops.reason.value_counts())          # 2 = land, 3 = left the domain, ...

The codes are listed in Stopping trajectories.

Transit times

The first time each particle enters a target region:

in_target = (df.lat > 60) & (df.lon.between(-30, 10))
first = df[in_target].groupby("id").dtm.min()
release = df.groupby("id").dtm.min()
transit_days = (first - release.loc[first.index]).dt.total_seconds() / 86400
transit_days.describe()

For a backward run, dtm decreases along each trajectory, so the same code gives the age of the water since it left the target region.

Cohorts of a repeat-seeded run

df = tracmass.load(env="gog_upwelling", columns=["id", "dtm", "x", "y", "t0"])
df["age_days"] = (df.groupby("id").dtm.transform("min") - df.dtm).abs().dt.days
df.groupby("t0").id.nunique()                    # particles per release

Positions at fixed times across runs

Positions are saved in the full-grid frame, so runs on different subsets can be compared directly:

a = tracmass.load(env="omi_small", filters={"window": 30})
b = tracmass.load(env="omi_large", filters={"window": 30})
merged = a.merge(b, on="id", suffixes=("_small", "_large"))

Geographic conversion without an adapter

tracmass.trajs.xy_to_lonlat converts positions on any grid with longitude and latitude arrays (see Coordinates):

import types
from tracmass.trajs import xy_to_lonlat

grid = types.SimpleNamespace(lon=lon2d, lat=lat2d)
lon, lat = xy_to_lonlat(df.x, df.y, grid)