scripts/train_predictor.py

#!/usr/bin/env python3
"""Train and validate the short-horizon density predictor.

The simulator is the data generator. Because it provides exact ground truth,
the model can be validated honestly: training and test use *disjoint seeds*,
and the report records the model's mean absolute error alongside the analytic
mass-balance baseline. If the model does not beat the baseline it is not used
at inference time.

Run:  python scripts/train_predictor.py [--quick]
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "backend"))

from flowtwin.config import MODEL_DIR, SETTINGS                    # noqa: E402
from flowtwin.prediction.features import (                          # noqa: E402
    FEATURE_NAMES,
    analytic_projection,
    build_feature_matrix,
)
from flowtwin.prediction.model import TrainingReport, fit_models    # noqa: E402
from flowtwin.simulation.engine import Simulator                    # noqa: E402
from flowtwin.venue import compile_venue, load_scenario             # noqa: E402

TRAIN_SEEDS = [42193, 1177, 90210, 5, 771]
TEST_SEEDS = [31337, 8080]
SCENARIOS = ["circuit_alpha_post_race", "barcelona_2022_egress", "circuit_alpha_arrival"]

#: Only sample every Nth step; consecutive steps are near-duplicates.
STEP_STRIDE = 4
#: Skip near-empty edges — they are trivially predictable and would dominate.
MIN_DENSITY = 0.05


def collect(scenario_id: str, seed: int, horizons: tuple[int, ...], max_steps: int
            ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Run one seeded scenario and return (features, targets, baseline)."""
    scenario = load_scenario(scenario_id)
    venue = compile_venue(scenario.venue_id)
    sim = Simulator(venue, scenario, SETTINGS, seed=seed)

    dt_s = sim.dt
    horizon_steps = [int(round(h / dt_s)) for h in horizons]
    max_h = max(horizon_steps)

    feats: list[np.ndarray] = []
    base: list[np.ndarray] = []
    density_track: list[np.ndarray] = []
    sample_at: list[int] = []

    steps = min(max_steps, int(scenario.duration_s / dt_s))
    for k in range(steps):
        sim.step()
        density_track.append(sim.state.edge_density.copy())
        if k % STEP_STRIDE == 0:
            feats.append(build_feature_matrix(sim))
            base.append(analytic_projection(sim, horizons))
            sample_at.append(k)
        if sim.is_complete and k > max_h:
            break

    if not sample_at:
        return (np.empty((0, len(FEATURE_NAMES)), np.float32),
                np.empty((len(horizons), 0), np.float32),
                np.empty((len(horizons), 0), np.float32))

    n_track = len(density_track)
    X_parts, Y_parts, B_parts = [], [], []
    for j, k in enumerate(sample_at):
        if k + max_h >= n_track:
            break
        now = density_track[k]
        keep = now >= MIN_DENSITY
        if not np.any(keep):
            continue
        X_parts.append(feats[j][keep])
        Y_parts.append(np.stack([density_track[k + hs][keep] for hs in horizon_steps]))
        B_parts.append(base[j][:, keep])

    if not X_parts:
        return (np.empty((0, len(FEATURE_NAMES)), np.float32),
                np.empty((len(horizons), 0), np.float32),
                np.empty((len(horizons), 0), np.float32))

    return (np.concatenate(X_parts, axis=0),
            np.concatenate(Y_parts, axis=1),
            np.concatenate(B_parts, axis=1))


def gather(seeds: list[int], horizons: tuple[int, ...], max_steps: int, label: str):
    Xs, Ys, Bs = [], [], []
    for scenario_id in SCENARIOS:
        for seed in seeds:
            X, Y, B = collect(scenario_id, seed, horizons, max_steps)
            if X.shape[0] == 0:
                continue
            Xs.append(X)
            Ys.append(Y)
            Bs.append(B)
            print(f"  [{label}] {scenario_id:<26} seed={seed:<8} rows={X.shape[0]:,}")
    return (np.concatenate(Xs, axis=0),
            np.concatenate(Ys, axis=1),
            np.concatenate(Bs, axis=1))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--quick", action="store_true",
                    help="fewer seeds and shorter runs, for a fast check")
    args = ap.parse_args()

    horizons = tuple(SETTINGS.prediction.horizons_s)
    train_seeds = TRAIN_SEEDS[:2] if args.quick else TRAIN_SEEDS
    test_seeds = TEST_SEEDS[:1] if args.quick else TEST_SEEDS
    max_steps = 900 if args.quick else 2600

    print(f"Horizons: {horizons}s   train seeds {train_seeds}   test seeds {test_seeds}")
    print("Generating training data ...")
    Xtr, Ytr, _ = gather(train_seeds, horizons, max_steps, "train")
    print("Generating held-out data ...")
    Xte, Yte, Bte = gather(test_seeds, horizons, max_steps, "test")
    print(f"train rows {Xtr.shape[0]:,}   test rows {Xte.shape[0]:,}")

    predictor, model_name = fit_models(Xtr, Ytr, horizons, seed=7)
    pred = predictor.predict(Xte)

    mae_model, mae_base, rmse_model, r2_model, improvement = {}, {}, {}, {}, {}
    for k, h in enumerate(horizons):
        err_m = np.abs(pred[k] - Yte[k])
        err_b = np.abs(Bte[k] - Yte[k])
        mae_model[str(h)] = float(err_m.mean())
        mae_base[str(h)] = float(err_b.mean())
        rmse_model[str(h)] = float(np.sqrt(((pred[k] - Yte[k]) ** 2).mean()))
        ss_res = float(((pred[k] - Yte[k]) ** 2).sum())
        ss_tot = float(((Yte[k] - Yte[k].mean()) ** 2).sum())
        r2_model[str(h)] = 1.0 - ss_res / max(ss_tot, 1e-9)
        improvement[str(h)] = 100.0 * (mae_base[str(h)] - mae_model[str(h)]) / max(mae_base[str(h)], 1e-9)

    print("\nhorizon   MAE model   MAE baseline   improvement      R²")
    for h in horizons:
        k = str(h)
        print(f"  +{h:>3}s      {mae_model[k]:.4f}        {mae_base[k]:.4f}"
              f"        {improvement[k]:+6.1f}%   {r2_model[k]:.3f}")

    if all(v <= 0 for v in improvement.values()):
        print("\nModel did not beat the analytic baseline. Not saving; inference "
              "will keep using the mass-balance projection.")
        return

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    predictor.save(SETTINGS.prediction.model_path)
    report = TrainingReport(
        horizons_s=list(horizons),
        n_train=int(Xtr.shape[0]),
        n_test=int(Xte.shape[0]),
        scenarios=SCENARIOS,
        train_seeds=train_seeds,
        test_seeds=test_seeds,
        model_name=model_name,
        mae_model=mae_model,
        mae_baseline=mae_base,
        rmse_model=rmse_model,
        r2_model=r2_model,
        improvement_pct=improvement,
        feature_names=list(FEATURE_NAMES),
        created_utc=dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
    )
    SETTINGS.prediction.metrics_path.write_text(report.to_json(), encoding="utf-8")
    print(f"\nSaved model  -> {SETTINGS.prediction.model_path}")
    print(f"Saved report -> {SETTINGS.prediction.metrics_path}")


if __name__ == "__main__":
    main()