eval/eval.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Per-satellite evaluation for the v2 (cleaned + solar) Toto TLE forecaster.

Same two metric families as v1, reported PER SATELLITE:
  (1) element RMSE (mean_motion, inclination, eccentricity; mean anomaly / RAAN /
      argp as circular error) vs truth and vs a persistence baseline;
  (2) SGP4 position error (km): predicted elements -> SGP4 -> TEME @ t0+Δ, vs
      truth; baseline = propagate the last observed TLE.

Framing: given n=context-days, forecast m=horizon-days. Solar channels are fed
as true context. For horizons <= patch_size we use a single forward pass
(decode_block_size=None), so no autoregressive feedback of predicted solar.

Run:
  python v2/eval/eval.py --ckpt v2/ckpt/toto_v2_Toto-2.0-4m.pt \
    --model Datadog/Toto-2.0-4m --years 2020 --split all \
    --sw-csv v2/data/SW-All.csv --context-days 64 --horizon-days 30 \
    --horizons 1 3 7 14 30 --max-samples 4000 --out-csv v2/eval_out/per_sat.csv
"""

from __future__ import annotations

import argparse
import csv
import math
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from tqdm import tqdm

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

from tle_dataset import (  # noqa: E402
    build_daily_series, elements_from_feat_aux, reconstruct_track, sat_split_of,
    N_CHANNELS, N_ORBITAL,
)
from tle_future_dataset import parse_date_to_unix  # noqa: E402
from toto2 import Toto2Model  # noqa: E402
from sgp4.api import Satrec, WGS72  # noqa: E402

_JD_UNIX_EPOCH = 2440587.5
_JD_SGP4_EPOCH = 2433281.5


def unix_to_jd(u):
    return u / 86400.0 + _JD_UNIX_EPOCH


def build_satrec(elem, epoch_unix, satnum=99999):
    sat = Satrec()
    sat.sgp4init(
        WGS72, "i", satnum, unix_to_jd(epoch_unix) - _JD_SGP4_EPOCH,
        float(elem["bstar"]), 0.0, 0.0, float(elem["eccentricity"]),
        math.radians(elem["argp_deg"]), math.radians(elem["inclination_deg"]),
        math.radians(elem["mean_anomaly_deg"]),
        elem["mean_motion_rev_per_day"] * 2.0 * math.pi / 1440.0,
        math.radians(elem["raan_deg"]),
    )
    return sat


def propagate(sat, target_unix):
    jd = unix_to_jd(target_unix)
    jd_i = math.floor(jd - 0.5) + 0.5
    e, r, v = sat.sgp4(jd_i, jd - jd_i)
    return None if e != 0 else np.asarray(r, dtype=np.float64)


ELEMS = ["mean_motion_rev_per_day", "inclination_deg", "eccentricity",
         "mean_anomaly_deg", "raan_deg", "argp_deg"]
CIRC = {"mean_anomaly_deg", "raan_deg", "argp_deg"}
SHORT = {"mean_motion_rev_per_day": "mm", "inclination_deg": "inc", "eccentricity": "ecc",
         "mean_anomaly_deg": "ma", "raan_deg": "raan", "argp_deg": "argp"}


def elem_err(pred, true, key):
    d = pred[key] - true[key]
    if key in CIRC:
        d = ((d + 180.0) % 360.0) - 180.0
    return abs(d)


def split_of(epoch, tu, vu):
    return "train" if epoch < tu else ("valid" if epoch < vu else "test")


def rmse(xs):
    return float(np.sqrt(np.mean(np.square(xs)))) if len(xs) else math.nan


@torch.no_grad()
def evaluate(model, series, device, patch_size, n_days, horizons, split, tu, vu,
             stride_days, per_sat_samples, max_eval_sats, batch_sats, recon="integrate",
             dump_k=0, split_mode="time"):
    model.eval()
    m_days = max(horizons)
    L = math.ceil(n_days / patch_size) * patch_size
    median_idx = model.output_head.knots.index(0.5)
    block = None if m_days <= patch_size else patch_size  # single forward for short horizons

    def empty():
        return {"pos_m": [], "pos_b": [], "e_m": defaultdict(list), "e_b": defaultdict(list)}
    per_sat = defaultdict(lambda: {h: empty() for h in horizons})

    # Group anchors PER satellite so each satellite gets several samples (real
    # per-satellite RMSE), instead of ~1 anchor each after a global subsample.
    by_sat = {}
    for norad, s in series.items():
        T = s.feats.shape[0]
        # satellite-level split: whole satellite belongs to one split (cm-tle-pred)
        if split != "all" and split_mode == "satellite" and sat_split_of(norad) != split:
            continue
        anchors = []
        for a in range(n_days - 1, T - m_days, stride_days):
            if not s.mask[a - n_days + 1: a + 1].all():
                continue
            if not s.mask[a + 1: a + m_days + 1].all():
                continue
            if split != "all" and split_mode == "time" \
                    and split_of(float(s.grid_epochs[a]), tu, vu) != split:
                continue
            anchors.append(a)
        if not anchors:
            continue
        if per_sat_samples and len(anchors) > per_sat_samples:
            pick = np.linspace(0, len(anchors) - 1, per_sat_samples).astype(int)
            anchors = [anchors[i] for i in sorted(set(pick))]
        by_sat[norad] = anchors

    sats = list(by_sat)
    if max_eval_sats and len(sats) > max_eval_sats:
        pick = np.linspace(0, len(sats) - 1, max_eval_sats).astype(int)
        sats = [sats[i] for i in sorted(set(pick))]
    jobs = [(n, a) for n in sats for a in by_sat[n]]
    print(f"[eval] {len(sats)} satellites x up to {per_sat_samples} samples = {len(jobs)} windows")

    n_fail = 0
    dumps = []   # concrete (truth | model | base) values for sanity-checking magnitudes
    for bstart in tqdm(range(0, len(jobs), batch_sats), desc="eval", unit="batch"):
        batch = jobs[bstart: bstart + batch_sats]
        tgt = torch.zeros(len(batch), N_CHANNELS, L, dtype=torch.float32)
        msk = torch.zeros(len(batch), N_CHANNELS, L, dtype=torch.bool)
        for bi, (norad, a) in enumerate(batch):
            ctx = series[norad].feats[a - n_days + 1: a + 1]
            tgt[bi, :, L - n_days:] = torch.from_numpy(ctx.T.copy())
            msk[bi, :, L - n_days:] = True
        sids = torch.zeros(len(batch), N_CHANNELS, dtype=torch.long)
        q = model.forecast({"target": tgt.to(device), "target_mask": msk.to(device),
                            "series_ids": sids.to(device)},
                           horizon=m_days, decode_block_size=block, has_missing_values=True)
        pred = q[median_idx].float().cpu().numpy()       # (B,C,m_days)

        for bi, (norad, a) in enumerate(batch):
            s = series[norad]
            anchor_aux, anchor_feat = s.aux[a], s.feats[a]
            t_a = float(s.grid_epochs[a])
            anchor_full = elements_from_feat_aux(anchor_feat, anchor_aux)
            track_m = reconstruct_track(anchor_aux, pred[bi].T[:m_days])
            # persistence baseline: zero drift, ecc/inc held at anchor, no angle drift
            base_orbital = np.array([0.0, 0.0, float(anchor_feat[2]), float(anchor_feat[3]), 0.0, 0.0])
            track_b = reconstruct_track(anchor_aux, np.tile(base_orbital, (m_days, 1)))
            base_sat = build_satrec(anchor_full, t_a)
            for h in horizons:
                j = a + h
                t_j = float(s.grid_epochs[j])
                true_el = elements_from_feat_aux(s.feats[j], s.aux[j])
                pred_el, base_el = track_m[h - 1], track_b[h - 1]
                rec = per_sat[norad][h]
                for k in ELEMS:
                    rec["e_m"][k].append(elem_err(pred_el, true_el, k))
                    rec["e_b"][k].append(elem_err(base_el, true_el, k))
                r_true = propagate(build_satrec(true_el, t_j), t_j)
                if recon == "sgp4":
                    # Drive SGP4 with the model's FORECAST, not just bstar: along-track
                    # comes from the model's predicted mean motion, interval-averaged over
                    # [t_a, t_j] so it already embodies the drag-driven decay. bstar is set
                    # to 0 so SGP4 does not re-apply that decay (double count). Phase stays
                    # analytic (SGP4), but now at the model's predicted rate.
                    mm_eff = float(np.mean([track_m[k]["mean_motion_rev_per_day"]
                                            for k in range(h)]))
                    model_sat = build_satrec(
                        {**anchor_full, "mean_motion_rev_per_day": mm_eff, "bstar": 0.0}, t_a)
                    r_model = propagate(model_sat, t_j)
                else:  # "integrate": daily trapezoidal phase reconstruction
                    r_model = propagate(build_satrec(pred_el, t_j), t_j)
                r_base = propagate(base_sat, t_j)
                if r_true is None or r_model is None or r_base is None:
                    n_fail += 1
                    continue
                pos_m = float(np.linalg.norm(r_model - r_true))
                pos_b = float(np.linalg.norm(r_base - r_true))
                rec["pos_m"].append(pos_m)
                rec["pos_b"].append(pos_b)
                if len(dumps) < dump_k:
                    bstar_model = float(pred_el["bstar"])   # model's forecast bstar
                    dumps.append({
                        "norad": norad, "h": h, "recon": recon,
                        "true": {k: float(true_el[k]) for k in ELEMS},
                        "model": {k: float(pred_el[k]) for k in ELEMS},
                        "base": {k: float(base_el[k]) for k in ELEMS},
                        "bstar_true": float(true_el["bstar"]),
                        "bstar_model": float(bstar_model),
                        "bstar_base": float(base_el["bstar"]),
                        "r_true": r_true.tolist(), "r_model": r_model.tolist(),
                        "r_base": r_base.tolist(), "pos_m": pos_m, "pos_b": pos_b,
                    })
    model.train()
    return per_sat, len(jobs), n_fail, dumps


def summarize(per_sat, horizons):
    rows = []
    for norad, hd in per_sat.items():
        row = {"norad": norad, "n": 0}
        for h in horizons:
            rec = hd[h]
            row["n"] = max(row["n"], len(rec["pos_m"]))
            row[f"posR_m_{h}"] = rmse(rec["pos_m"]); row[f"posR_b_{h}"] = rmse(rec["pos_b"])
            for k in ELEMS:
                row[f"{SHORT[k]}R_m_{h}"] = rmse(rec["e_m"][k])
                row[f"{SHORT[k]}R_b_{h}"] = rmse(rec["e_b"][k])
        rows.append(row)
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default=None)
    ap.add_argument("--model", default="Datadog/Toto-2.0-4m")
    ap.add_argument("--input-dir", default="/home/irteam/data-vol1/models/OrbitGPT/data/TLEs")
    ap.add_argument("--cache-dir", default="/home/irteam/data-vol1/models/OrbitGPT/v2/cache")
    ap.add_argument("--cache-file", default=None,
                    help="explicit prebuilt cache npz (e.g. the full 2005-2024 superset); "
                         "filters by --split")
    ap.add_argument("--sw-csv", default="/home/irteam/data-vol1/models/OrbitGPT/v2/data/SW-All.csv")
    ap.add_argument("--years", type=int, nargs="+", default=[2020])
    ap.add_argument("--split", default="all")
    ap.add_argument("--split-mode", default="time", choices=["time", "satellite"],
                    help="time = epoch cutoffs; satellite = cm-tle-pred 70/15/15 by NORAD")
    ap.add_argument("--no-clean", action="store_true")
    ap.add_argument("--no-leo", action="store_true", help="disable LEO-only filter")
    ap.add_argument("--train-until", default="2022-01-01")
    ap.add_argument("--valid-until", default="2023-01-01")
    ap.add_argument("--context-days", type=int, default=64)
    ap.add_argument("--horizon-days", type=int, default=30)
    ap.add_argument("--horizons", type=int, nargs="+", default=[1, 3, 7, 14, 30])
    ap.add_argument("--stride-days", type=int, default=15)
    ap.add_argument("--per-sat-samples", type=int, default=8,
                    help="max evaluation anchors per satellite (gives n>1 per-sat RMSE)")
    ap.add_argument("--max-eval-sats", type=int, default=1500,
                    help="max number of satellites to evaluate")
    ap.add_argument("--max-satellites", type=int, default=None)
    ap.add_argument("--batch-sats", type=int, default=64)
    ap.add_argument("--recon", default="sgp4", choices=["integrate", "sgp4"],
                    help="model position reconstruction: 'integrate' (daily trapezoidal "
                         "phase) or 'sgp4' (SGP4 analytic phase + model bstar drag correction)")
    ap.add_argument("--dump-samples", type=int, default=8,
                    help="print this many concrete (truth|model|base) rows for sanity-checking")
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--out-csv", default="/home/irteam/data-vol1/models/OrbitGPT/v2/eval_out/per_sat.csv")
    ap.add_argument("--show", type=int, default=12)
    args = ap.parse_args()

    horizons = [h for h in args.horizons if h <= args.horizon_days]
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    model = Toto2Model.from_pretrained(args.model).to(device)
    patch_size = model.config.patch_size
    if args.ckpt:
        sd = torch.load(args.ckpt, map_location=device)
        model.load_state_dict(sd["model"] if "model" in sd else sd)
        print(f"[eval] loaded checkpoint {args.ckpt}")
    else:
        print("[eval] zero-shot pretrained weights")

    series = build_daily_series(
        args.input_dir, years=args.years, cache_dir=args.cache_dir, cache_file=args.cache_file,
        sw_csv=args.sw_csv, clean=not args.no_clean, leo_only=not args.no_leo,
        min_grid_points=args.context_days + args.horizon_days, verbose=True,
    )
    if args.max_satellites is not None:
        keep = sorted(series.keys())[: args.max_satellites]
        series = {k: series[k] for k in keep}

    per_sat, n_jobs, n_fail, dumps = evaluate(
        model, series, device, patch_size, args.context_days, horizons, args.split,
        parse_date_to_unix(args.train_until), parse_date_to_unix(args.valid_until),
        args.stride_days, args.per_sat_samples, args.max_eval_sats, args.batch_sats,
        recon=args.recon, dump_k=args.dump_samples, split_mode=args.split_mode,
    )
    print(f"[eval] reconstruction mode: {args.recon}")
    rows = summarize(per_sat, horizons)
    rows.sort(key=lambda r: r["norad"])
    print(f"\n[eval] n={args.context_days}d ctx -> m={args.horizon_days}d | "
          f"samples={n_jobs} satellites={len(rows)} sgp4_fail={n_fail}")

    hdr = ["norad", "n"]
    for h in horizons:
        hdr += [f"posR_m_{h}", f"posR_b_{h}"]
        for k in ELEMS:
            hdr += [f"{SHORT[k]}R_m_{h}", f"{SHORT[k]}R_b_{h}"]
    Path(args.out_csv).parent.mkdir(parents=True, exist_ok=True)
    with open(args.out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=hdr); w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k, "") for k in hdr})
    print(f"[eval] per-satellite CSV ({len(hdr)} cols) -> {args.out_csv}")

    print("\nper-satellite position RMSE model/base (km):")
    print(f"{'norad':>7} {'n':>4}  " + "  ".join(f"{f'{h}d':>15}" for h in horizons))
    for r in rows[: args.show]:
        print(f"{r['norad']:>7} {r['n']:>4}  " +
              "  ".join(f"{r[f'posR_m_{h}']:>6.0f}/{r[f'posR_b_{h}']:<8.1f}" for h in horizons))

    hmax = max(horizons)
    print(f"\nper-satellite element RMSE @ {hmax}d  (model | baseline):")
    print(f"{'norad':>7} {'n':>4}  {'mm(rev/d)':>20}  {'ma(deg)':>16}  {'inc(deg)':>16}  {'ecc':>20}")
    for r in rows[: args.show]:
        print(f"{r['norad']:>7} {r['n']:>4}  "
              f"{r[f'mmR_m_{hmax}']:>9.6f}|{r[f'mmR_b_{hmax}']:<10.6f}  "
              f"{r[f'maR_m_{hmax}']:>7.2f}|{r[f'maR_b_{hmax}']:<8.2f}  "
              f"{r[f'incR_m_{hmax}']:>7.4f}|{r[f'incR_b_{hmax}']:<8.4f}  "
              f"{r[f'eccR_m_{hmax}']:>9.6f}|{r[f'eccR_b_{hmax}']:<10.6f}")

    if dumps:
        print("\nground-truth check — actual values (true | model | base):")
        print("  [recon=sgp4 → position uses ONLY bstar; printed mm/ma/inc/ecc are the "
              "model's element forecast, NOT what drives r_model]")
        for d in dumps:
            t, m, b = d["true"], d["model"], d["base"]
            rt, rm, rb = d["r_true"], d["r_model"], d["r_base"]
            print(f"  norad {d['norad']} @{d['h']}d  [recon={d['recon']}]")
            print(f"    mm(rev/d) true={t['mean_motion_rev_per_day']:.7f}  "
                  f"model={m['mean_motion_rev_per_day']:.7f}  base={b['mean_motion_rev_per_day']:.7f}")
            print(f"    ma(deg)   true={t['mean_anomaly_deg']:8.3f}  "
                  f"model={m['mean_anomaly_deg']:8.3f}  base={b['mean_anomaly_deg']:8.3f}")
            print(f"    inc(deg)  true={t['inclination_deg']:8.4f}  "
                  f"model={m['inclination_deg']:8.4f}  base={b['inclination_deg']:8.4f}")
            print(f"    ecc       true={t['eccentricity']:.7f}  "
                  f"model={m['eccentricity']:.7f}  base={b['eccentricity']:.7f}")
            print(f"    bstar     true={d['bstar_true']:.6e}  "
                  f"model={d['bstar_model']:.6e}  base={d['bstar_base']:.6e}")
            print(f"    |r_true|={math.sqrt(sum(x*x for x in rt)):.1f}km   "
                  f"pos_err: model={d['pos_m']:.1f}km  base={d['pos_b']:.1f}km")
            print(f"      r_true =[{rt[0]:9.1f},{rt[1]:9.1f},{rt[2]:9.1f}]")
            print(f"      r_model=[{rm[0]:9.1f},{rm[1]:9.1f},{rm[2]:9.1f}]")
            print(f"      r_base =[{rb[0]:9.1f},{rb[1]:9.1f},{rb[2]:9.1f}]")

    def med(key):
        v = np.array([r[key] for r in rows if not math.isnan(r.get(key, math.nan))])
        return np.median(v) if len(v) else math.nan
    def sat_win(mk, bk):
        m = np.array([r[mk] for r in rows]); b = np.array([r[bk] for r in rows])
        ok = ~(np.isnan(m) | np.isnan(b))
        return float(np.mean(m[ok] < b[ok]) * 100.0) if ok.any() else math.nan

    print("\naggregate over satellites (median per-sat RMSE) | model vs baseline:")
    print(f"{'horizon':>8} {'posKm_m':>9} {'posKm_b':>9} {'satwin%':>7} | "
          f"{'mm_m':>8} {'mm_b':>8} {'ma_m':>8} {'ma_b':>8}")
    for h in horizons:
        print(f"{h:>6}d {med(f'posR_m_{h}'):>9.1f} {med(f'posR_b_{h}'):>9.1f} "
              f"{sat_win(f'posR_m_{h}', f'posR_b_{h}'):>7.1f} | "
              f"{med(f'mmR_m_{h}'):>8.5f} {med(f'mmR_b_{h}'):>8.5f} "
              f"{med(f'maR_m_{h}'):>8.2f} {med(f'maR_b_{h}'):>8.2f}")


if __name__ == "__main__":
    main()