#!/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()