quickstart.py

#!/usr/bin/env python
"""MARIS-Forecast quickstart.

Loads one jurisdictional leaf, evaluates a constant-velocity dead-reckoning
baseline with the shipped ADE/FDE metrics, and prints the aligned context
tensor shapes.  Point ``--root`` at a local copy of the dataset downloaded
from Zenodo (DOI 10.5281/zenodo.21224009) or the Hugging Face mirror.

    python examples/quickstart.py --root /data/maris --track A --region dma

The leaf directory can also be given directly with ``--leaf`` (handy when a
single leaf has been extracted on its own).
"""
import argparse
import os
import sys

import numpy as np

# allow running without `pip install -e .`
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from maris_forecast import loader, metrics  # noqa: E402


def dead_reckoning(hist, n_future):
    """Constant-velocity extrapolation from the last two observed points."""
    v = hist[:, -1, :] - hist[:, -2, :]                 # (N, 2) per-step velocity
    steps = np.arange(1, n_future + 1)[None, :, None]   # (1, T, 1)
    return hist[:, -1:, :] + steps * v[:, None, :]      # (N, T, 2)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--root", help="root of a downloaded MARIS-Forecast release")
    ap.add_argument("--leaf", help="path to a single leaf directory (overrides --root)")
    ap.add_argument("--track", default="A", choices=["A", "B"])
    ap.add_argument("--region", default="dma", choices=list(loader.schema.REGIONS))
    ap.add_argument("--split", default="test")
    ap.add_argument("--no-context", action="store_true",
                    help="skip loading environment/social tensors")
    args = ap.parse_args()

    leaf = args.leaf or loader.leaf_dir(args.root, args.track, args.region)
    print(f"leaf: {leaf}")

    df = loader.load_split(leaf, args.split, consistent_only=True)
    print(f"loaded {len(df):,} OSM-consistent {args.region.upper()} "
          f"Track {args.track} {args.split} samples")

    hist, fut = loader.stack_trajectories(df)
    pred = dead_reckoning(hist, fut.shape[1])
    print(f"  history {hist.shape}  future {fut.shape}")
    print(f"  dead-reckoning   ADE = {metrics.ade(pred, fut):6.1f} m   "
          f"FDE = {metrics.fde(pred, fut):6.1f} m")

    if not args.no_context:
        env = loader.load_environment(leaf, args.split, align_to=df.head(64))
        print(f"  env masks {tuple(np.asarray(env['masks']).shape)}  "
              f"shore-SDF {tuple(np.asarray(env['signed_dist_shore']).shape)}")
        soc = loader.load_social(leaf, args.split)
        nc = soc["neighbor_count_used"].fillna(0).astype(int)
        print(f"  social: {(nc > 0).mean() * 100:.1f}% of samples have "
              f"≥1 neighbour within 3 km")


if __name__ == "__main__":
    main()