Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch | |
| import dataclasses | |
| import json | |
| import math | |
| import os | |
| import time | |
| from datetime import datetime, timedelta, timezone | |
| from pathlib import Path | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import plotly.graph_objects as go | |
| import requests | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from plotly.subplots import make_subplots | |
| from sgp4.api import WGS72, Satrec | |
| from toto2 import Toto2Model | |
| from toto2.configuration import Toto2ModelConfig | |
| # --------------------------------------------------------------------------- | |
| # Constants (ported 1:1 from postcn/OrbitFM utils/tle_dataset.py) | |
| # --------------------------------------------------------------------------- | |
| MODEL_ID = "postcn/OrbitFM" | |
| BASE_ID = os.environ.get("ORBITFM_BASE", "Datadog/Toto-2.0-2.5B") | |
| LOCAL_TEST = os.environ.get("ORBITFM_LOCAL_TEST") == "1" # tiny base, no ckpt: pipeline smoke test only | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| CKPT_FILE = "ckpt/toto_v2_Toto-2.0-2.5B-001.pt" | |
| N_ORBITAL, N_PHYS, N_SOLAR = 6, 16, 3 | |
| N_CHANNELS = N_ORBITAL + N_PHYS + N_SOLAR # 25 | |
| AUX_MA, AUX_RAAN, AUX_ARGP, AUX_MM, AUX_BSTAR = 0, 1, 2, 3, 4 | |
| SECONDS_PER_DAY = 86400.0 | |
| MU_EARTH = 398600.4418 | |
| R_EARTH = 6378.137 | |
| LEO_MIN_MEAN_MOTION = 11.25 | |
| PATCH = 32 | |
| _JD_UNIX_EPOCH = 2440587.5 | |
| _JD_SGP4_EPOCH = 2433281.5 | |
| # --------------------------------------------------------------------------- | |
| # Model: Toto-2.0-2.5B architecture + OrbitFM fine-tuned state dict | |
| # --------------------------------------------------------------------------- | |
| cfg_raw = json.loads(Path(hf_hub_download(BASE_ID, "config.json")).read_text()) | |
| known = {f.name for f in dataclasses.fields(Toto2ModelConfig)} | |
| cfg = Toto2ModelConfig(**{k: v for k, v in cfg_raw.items() if k in known}) | |
| model = Toto2Model(cfg) | |
| if not LOCAL_TEST: | |
| ckpt_path = hf_hub_download(MODEL_ID, CKPT_FILE) | |
| sd = torch.load(ckpt_path, map_location="cpu", mmap=True, weights_only=False) | |
| sd = sd["model"] if isinstance(sd, dict) and "model" in sd else sd | |
| sd = {k.replace("_orig_mod.", "").removeprefix("module."): v for k, v in sd.items()} | |
| missing, unexpected = model.load_state_dict(sd, strict=False) | |
| print(f"[orbitfm] loaded fine-tuned checkpoint: missing={len(missing)} unexpected={len(unexpected)}") | |
| if missing or unexpected: | |
| print(" missing:", missing[:10], "unexpected:", unexpected[:10]) | |
| del sd | |
| model = model.to(DEVICE).eval() | |
| MEDIAN_IDX = model.output_head.knots.index(0.5) | |
| Q10_IDX = model.output_head.knots.index(0.1) | |
| Q90_IDX = model.output_head.knots.index(0.9) | |
| # --------------------------------------------------------------------------- | |
| # Space weather (CelesTrak SW-All.csv): fresh copy if reachable, else the | |
| # copy shipped in the model repo. | |
| # --------------------------------------------------------------------------- | |
| def _load_sw(): | |
| path = Path("SW-All.csv") | |
| try: | |
| r = requests.get("https://celestrak.org/SpaceData/SW-All.csv", timeout=20) | |
| r.raise_for_status() | |
| path.write_bytes(r.content) | |
| src = "celestrak.org (live)" | |
| except Exception as e: # noqa: BLE001 | |
| print(f"[sw] CelesTrak fetch failed ({e}); using model-repo copy") | |
| path = Path(hf_hub_download(MODEL_ID, "data/SW-All.csv")) | |
| src = "postcn/OrbitFM data/SW-All.csv" | |
| df = pd.read_csv(path) | |
| cols = {c.upper(): c for c in df.columns} | |
| def col(*names): | |
| for nm in names: | |
| if nm in cols: | |
| return df[cols[nm]] | |
| raise KeyError(names) | |
| dates = pd.to_datetime(col("DATE")) | |
| tab = pd.DataFrame({ | |
| "f107": pd.to_numeric(col("F10.7_OBS", "F10.7_ADJ"), errors="coerce"), | |
| "f107_81": pd.to_numeric(col("F10.7_OBS_CENTER81", "F10.7_ADJ_CENTER81", "F10.7_OBS_LAST81"), errors="coerce"), | |
| "ap": pd.to_numeric(col("AP_AVG"), errors="coerce"), | |
| }).ffill().bfill() | |
| day_unix = np.array([d.replace(tzinfo=timezone.utc).timestamp() for d in np.array(dates.dt.to_pydatetime())]) | |
| order = np.argsort(day_unix) | |
| print(f"[sw] {len(day_unix)} daily records from {src}, last {dates.max().date()}") | |
| return day_unix[order], tab.to_numpy(dtype=np.float32)[order], src | |
| SW_DAYS, SW_TABLE, SW_SRC = _load_sw() | |
| def sw_for_epochs(epochs): | |
| idx = np.clip(np.searchsorted(SW_DAYS, epochs, side="right") - 1, 0, len(SW_DAYS) - 1) | |
| return SW_TABLE[idx] | |
| # --------------------------------------------------------------------------- | |
| # TLE parsing + cleaning (port of utils/tle_clean.py) | |
| # --------------------------------------------------------------------------- | |
| class TLERecord: | |
| norad_id: int | |
| epoch_unix: float | |
| bstar: float | |
| inclination_deg: float | |
| raan_deg: float | |
| eccentricity: float | |
| argp_deg: float | |
| mean_anomaly_deg: float | |
| mean_motion_rev_per_day: float | |
| def signed_log1p(x): | |
| return math.copysign(math.log1p(abs(x)), x) | |
| def signed_expm1(x): | |
| return math.copysign(math.expm1(abs(x)), x) | |
| def _tle_exp(field): | |
| s = field.strip().replace(" ", "") | |
| if not s: | |
| return 0.0 | |
| sign = -1.0 if s[0] == "-" else 1.0 | |
| s = s.lstrip("+-") | |
| mant, ex = s[:-2], s[-2:] | |
| if ex[0] not in "+-": | |
| return sign * float("0." + s) | |
| return sign * float("0." + mant) * 10.0 ** int(ex) | |
| def parse_tle_text(text): | |
| lines = [ln.rstrip() for ln in text.splitlines() if ln.strip()] | |
| recs, name = [], None | |
| i = 0 | |
| while i < len(lines) - 1: | |
| l1, l2 = lines[i].strip(), lines[i + 1].strip() | |
| if l1.startswith("1 ") and l2.startswith("2 ") and len(l1) >= 63 and len(l2) >= 63: | |
| try: | |
| yy = int(l1[18:20]) | |
| year = 2000 + yy if yy < 57 else 1900 + yy | |
| doy = float(l1[20:32]) | |
| ep = datetime(year, 1, 1, tzinfo=timezone.utc) + timedelta(days=doy - 1) | |
| rec = TLERecord( | |
| norad_id=int(l1[2:7]), | |
| epoch_unix=ep.timestamp(), | |
| bstar=_tle_exp(l1[53:61]), | |
| inclination_deg=float(l2[8:16]), | |
| raan_deg=float(l2[17:25]), | |
| eccentricity=float("0." + l2[26:33].strip()), | |
| argp_deg=float(l2[34:42]), | |
| mean_anomaly_deg=float(l2[43:51]), | |
| mean_motion_rev_per_day=float(l2[52:63]), | |
| ) | |
| if (0.0 <= rec.eccentricity < 1.0 and 0.0 < rec.inclination_deg < 180.0 | |
| and 0.1 < rec.mean_motion_rev_per_day < 20.0): | |
| recs.append(rec) | |
| except ValueError: | |
| pass | |
| i += 2 | |
| else: | |
| if not l1.startswith(("1 ", "2 ")): | |
| name = l1 | |
| i += 1 | |
| return recs, name | |
| def cumulative_mean_anomaly(recs): | |
| phi = np.empty(len(recs)) | |
| phi[0] = recs[0].mean_anomaly_deg | |
| for i in range(1, len(recs)): | |
| dt = (recs[i].epoch_unix - recs[i - 1].epoch_unix) / SECONDS_PER_DAY | |
| n_avg = 0.5 * (recs[i].mean_motion_rev_per_day + recs[i - 1].mean_motion_rev_per_day) | |
| pred = phi[i - 1] + n_avg * dt * 360.0 | |
| m_i = recs[i].mean_anomaly_deg | |
| phi[i] = 360.0 * round((pred - m_i) / 360.0) + m_i | |
| return phi | |
| def _resid(t, x): | |
| r = np.zeros_like(x, dtype=np.float64) | |
| if len(x) < 3: | |
| return r | |
| t0, t1, t2 = t[:-2], t[1:-1], t[2:] | |
| w = (t1 - t0) / np.where((t2 - t0) == 0, 1.0, t2 - t0) | |
| r[1:-1] = x[1:-1] - (x[:-2] + (x[2:] - x[:-2]) * w) | |
| return r | |
| def _mad(v): | |
| v = v[np.isfinite(v)] | |
| return 0.0 if v.size == 0 else float(1.4826 * np.median(np.abs(v - np.median(v)))) | |
| def clean_records(recs, k_mad=6.0, max_passes=2): | |
| recs = sorted(recs, key=lambda r: r.epoch_unix) | |
| removed = 0 | |
| for _ in range(max_passes): | |
| n = len(recs) | |
| if n < 5: | |
| break | |
| t = np.array([r.epoch_unix for r in recs]) | |
| series = [ | |
| np.array([r.mean_motion_rev_per_day for r in recs]), | |
| np.array([r.inclination_deg for r in recs]), | |
| np.array([r.eccentricity for r in recs]), | |
| np.array([signed_log1p(r.bstar) for r in recs]), | |
| cumulative_mean_anomaly(recs), | |
| np.degrees(np.unwrap(np.radians([r.raan_deg for r in recs]))), | |
| np.degrees(np.unwrap(np.radians([r.argp_deg for r in recs]))), | |
| ] | |
| flag = np.zeros(n, dtype=bool) | |
| for s in series: | |
| res = _resid(t, s) | |
| sc = _mad(res[1:-1]) | |
| if sc > 0: | |
| flag |= np.abs(res) > k_mad * sc | |
| flag[0] = flag[-1] = False | |
| if not flag.any(): | |
| break | |
| recs = [r for r, bad in zip(recs, flag) if not bad] | |
| removed += int(flag.sum()) | |
| return recs, removed | |
| # --------------------------------------------------------------------------- | |
| # Daily grid + 25-channel features (port of utils/tle_dataset.py::_build_one) | |
| # --------------------------------------------------------------------------- | |
| def _kepler_rv(a, e, inc_r, raan_r, argp_r, M_r): | |
| E = np.array(M_r, dtype=np.float64, copy=True) | |
| for _ in range(12): | |
| E = E - (E - e * np.sin(E) - M_r) / np.maximum(1.0 - e * np.cos(E), 1e-9) | |
| cosE, sinE = np.cos(E), np.sin(E) | |
| r = a * (1.0 - e * cosE) | |
| nu = np.arctan2(np.sqrt(np.maximum(1.0 - e * e, 0.0)) * sinE, cosE - e) | |
| cosnu, sinnu = np.cos(nu), np.sin(nu) | |
| p = np.maximum(a * (1.0 - e * e), 1e-6) | |
| rp_x, rp_y = r * cosnu, r * sinnu | |
| vfac = np.sqrt(MU_EARTH / p) | |
| vp_x, vp_y = -vfac * sinnu, vfac * (e + cosnu) | |
| co, so = np.cos(raan_r), np.sin(raan_r) | |
| ci, si = np.cos(inc_r), np.sin(inc_r) | |
| cw, sw = np.cos(argp_r), np.sin(argp_r) | |
| R11 = co * cw - so * sw * ci | |
| R12 = -co * sw - so * cw * ci | |
| R21 = so * cw + co * sw * ci | |
| R22 = -so * sw + co * cw * ci | |
| R31, R32 = sw * si, cw * si | |
| return (R11 * rp_x + R12 * rp_y, R21 * rp_x + R22 * rp_y, R31 * rp_x + R32 * rp_y, | |
| R11 * vp_x + R12 * vp_y, R21 * vp_x + R22 * vp_y, R31 * vp_x + R32 * vp_y) | |
| def physics_features(g_mm, g_ecc, g_inc, g_raan, g_argp, g_ma): | |
| mm = np.maximum(np.asarray(g_mm, dtype=np.float64), 1e-6) | |
| n = mm * 2.0 * math.pi / SECONDS_PER_DAY | |
| a = (MU_EARTH / n ** 2) ** (1.0 / 3.0) | |
| e = np.clip(np.asarray(g_ecc, dtype=np.float64), 0.0, 0.999999) | |
| inc_r = np.radians(g_inc) | |
| raan_r = np.radians(np.asarray(g_raan) % 360.0) | |
| argp_r = np.radians(np.asarray(g_argp) % 360.0) | |
| ma_r = np.radians(np.asarray(g_ma) % 360.0) | |
| rx, ry, rz, vx, vy, vz = _kepler_rv(a, e, inc_r, raan_r, argp_r, ma_r) | |
| return np.stack([rx, ry, rz, vx, vy, vz, a, 1440.0 / mm, | |
| a * (1.0 + e) - R_EARTH, a * (1.0 - e) - R_EARTH, | |
| np.cos(ma_r), np.sin(ma_r), np.cos(raan_r), np.sin(raan_r), | |
| np.cos(argp_r), np.sin(argp_r)], axis=1) | |
| def build_series(recs): | |
| """Returns dict(grid, feats (T,25), mask (T,), aux (T,5)) or None.""" | |
| if len(recs) < 2: | |
| return None | |
| recs = sorted(recs, key=lambda r: r.epoch_unix) | |
| ep = np.array([r.epoch_unix for r in recs]) | |
| bstar_log = np.array([signed_log1p(r.bstar) for r in recs]) | |
| mm = np.array([r.mean_motion_rev_per_day for r in recs]) | |
| ecc = np.array([r.eccentricity for r in recs]) | |
| inc = np.array([r.inclination_deg for r in recs]) | |
| raan_uw = np.degrees(np.unwrap(np.radians([r.raan_deg for r in recs]))) | |
| argp_uw = np.degrees(np.unwrap(np.radians([r.argp_deg for r in recs]))) | |
| phiM = cumulative_mean_anomaly(recs) | |
| step = SECONDS_PER_DAY | |
| grid = np.arange(math.floor(ep[0] / step) * step, math.ceil(ep[-1] / step) * step + 0.5 * step, step) | |
| if grid.size < 2: | |
| return None | |
| itp = lambda y: np.interp(grid, ep, y) # noqa: E731 | |
| g_raan, g_argp, g_mm, g_bstar = itp(raan_uw), itp(argp_uw), itp(mm), itp(bstar_log) | |
| g_ecc, g_inc, g_ma = itp(ecc), itp(inc), itp(phiM) | |
| orbital = np.stack([np.gradient(g_bstar), np.gradient(g_mm), g_ecc, g_inc, | |
| np.gradient(g_raan), np.gradient(g_argp)], axis=1) | |
| phys = physics_features(g_mm, g_ecc, g_inc, g_raan, g_argp, g_ma) | |
| solar = sw_for_epochs(grid) | |
| feats = np.concatenate([orbital, phys, solar], axis=1).astype(np.float32) | |
| aux = np.stack([g_ma % 360.0, g_raan % 360.0, g_argp % 360.0, g_mm, g_bstar], axis=1).astype(np.float32) | |
| idx = np.clip(np.searchsorted(ep, grid), 1, len(ep) - 1) | |
| gap = np.minimum(np.abs(grid - ep[idx]), np.abs(grid - ep[idx - 1])) | |
| return {"grid": grid, "feats": feats, "mask": gap <= 2.0 * step, "aux": aux} | |
| def elements_from_feat_aux(f, a): | |
| return { | |
| "bstar": signed_expm1(float(a[AUX_BSTAR])), | |
| "mean_motion_rev_per_day": float(max(a[AUX_MM], 1e-6)), | |
| "eccentricity": float(min(max(f[2], 0.0), 0.9999999)), | |
| "inclination_deg": float(min(max(f[3], 0.0), 180.0)), | |
| "mean_anomaly_deg": float(a[AUX_MA] % 360.0), | |
| "raan_deg": float(a[AUX_RAAN] % 360.0), | |
| "argp_deg": float(a[AUX_ARGP] % 360.0), | |
| } | |
| def reconstruct_track(anchor_aux, pred): | |
| a = np.asarray(anchor_aux, dtype=np.float64) | |
| ma, raan, argp = a[AUX_MA], a[AUX_RAAN], a[AUX_ARGP] | |
| mm, bs = float(a[AUX_MM]), float(a[AUX_BSTAR]) | |
| prev = mm | |
| out = [] | |
| for h in range(pred.shape[0]): | |
| bs += float(pred[h, 0]) | |
| mm += float(pred[h, 1]) | |
| mm_h = max(mm, 1e-6) | |
| ma += 360.0 * 0.5 * (prev + mm_h) | |
| prev = mm_h | |
| raan += float(pred[h, 4]) | |
| argp += float(pred[h, 5]) | |
| out.append({ | |
| "bstar": signed_expm1(bs), "mean_motion_rev_per_day": mm_h, | |
| "eccentricity": float(min(max(pred[h, 2], 0.0), 0.9999999)), | |
| "inclination_deg": float(min(max(pred[h, 3], 0.0), 180.0)), | |
| "mean_anomaly_deg": ma % 360.0, "raan_deg": raan % 360.0, "argp_deg": argp % 360.0, | |
| }) | |
| return out | |
| # --------------------------------------------------------------------------- | |
| # SGP4 helpers (port of eval/eval.py) | |
| # --------------------------------------------------------------------------- | |
| def unix_to_jd(u): | |
| return u / 86400.0 + _JD_UNIX_EPOCH | |
| def build_satrec(el, epoch_unix): | |
| sat = Satrec() | |
| sat.sgp4init(WGS72, "i", 99999, unix_to_jd(epoch_unix) - _JD_SGP4_EPOCH, | |
| float(el["bstar"]), 0.0, 0.0, float(el["eccentricity"]), | |
| math.radians(el["argp_deg"]), math.radians(el["inclination_deg"]), | |
| math.radians(el["mean_anomaly_deg"]), | |
| el["mean_motion_rev_per_day"] * 2.0 * math.pi / 1440.0, math.radians(el["raan_deg"])) | |
| return sat | |
| def propagate(sat, t): | |
| jd = unix_to_jd(t) | |
| jd_i = math.floor(jd - 0.5) + 0.5 | |
| e, r, _ = sat.sgp4(jd_i, jd - jd_i) | |
| return None if e != 0 else np.asarray(r, dtype=np.float64) | |
| def mm_to_alt(mm): | |
| n = np.maximum(np.asarray(mm, dtype=np.float64), 1e-6) * 2.0 * math.pi / SECONDS_PER_DAY | |
| return (MU_EARTH / n ** 2) ** (1.0 / 3.0) - R_EARTH | |
| def _date(u): | |
| return datetime.fromtimestamp(float(u), tz=timezone.utc).strftime("%Y-%m-%d") | |
| def _dt(u): | |
| return datetime.fromtimestamp(float(u), tz=timezone.utc) | |
| # --------------------------------------------------------------------------- | |
| # GPU step: one forward pass of OrbitFM | |
| # --------------------------------------------------------------------------- | |
| def _forecast_gpu(target: np.ndarray, mask: np.ndarray, horizon: int): | |
| tgt = torch.from_numpy(target).unsqueeze(0).to(DEVICE) | |
| msk = torch.from_numpy(mask).unsqueeze(0).to(DEVICE) | |
| sids = torch.zeros(1, N_CHANNELS, dtype=torch.long, device=DEVICE) | |
| block = None if horizon <= PATCH else PATCH | |
| with torch.no_grad(): | |
| q = model.forecast({"target": tgt, "target_mask": msk, "series_ids": sids}, | |
| horizon=horizon, decode_block_size=block, has_missing_values=True) | |
| q = q.float().cpu().numpy()[:, 0] # (9, C, H) | |
| return q[MEDIAN_IDX], q[Q10_IDX], q[Q90_IDX] | |
| # --------------------------------------------------------------------------- | |
| # Main entry | |
| # --------------------------------------------------------------------------- | |
| def forecast_orbit(tle_file: str, anchor_date: str = "", horizon_days: int = 30, | |
| context_days: int = 64, tle_text: str = ""): | |
| """Forecast a LEO satellite's orbital elements and positions with OrbitFM. | |
| Args: | |
| tle_file: path to a text file with the satellite's TLE history (2- or 3-line format, | |
| ideally at least ~70 days of TLEs, one object only). | |
| anchor_date: forecast start date (YYYY-MM-DD, UTC). Empty = latest date for which the | |
| full horizon can be checked against the real TLEs that followed (backtest). | |
| A date at/after the last TLE gives a pure forecast with no ground truth. | |
| horizon_days: forecast horizon in days (1-30). | |
| context_days: days of daily-grid history fed to the model (64 matches the paper's eval). | |
| tle_text: optional pasted TLE text, used instead of the file when non-empty. | |
| Returns: | |
| summary markdown, position-error plot, element plot, per-day forecast table, CSV path. | |
| """ | |
| t_start = time.perf_counter() | |
| if tle_text and tle_text.strip(): | |
| text = tle_text | |
| elif tle_file: | |
| text = Path(tle_file).read_text(errors="ignore") | |
| else: | |
| raise gr.Error("Upload a TLE history file or paste TLE text.") | |
| recs, name = parse_tle_text(text) | |
| if len(recs) < 10: | |
| raise gr.Error(f"Found only {len(recs)} valid TLEs - need a history of at least ~70 days.") | |
| ids = pd.Series([r.norad_id for r in recs]).value_counts() | |
| norad = int(ids.index[0]) | |
| if len(ids) > 1: | |
| gr.Warning(f"Multiple objects found; using NORAD {norad} (most TLEs).") | |
| recs = [r for r in recs if r.norad_id == norad] | |
| recs = sorted(recs, key=lambda r: r.epoch_unix) | |
| dedup, seen = [], set() | |
| for r in recs: | |
| k = round(r.epoch_unix, 3) | |
| if k not in seen: | |
| seen.add(k) | |
| dedup.append(r) | |
| recs, n_removed = clean_records(dedup) | |
| med_mm = float(np.median([r.mean_motion_rev_per_day for r in recs])) | |
| if med_mm < LEO_MIN_MEAN_MOTION: | |
| gr.Warning(f"Median mean motion {med_mm:.2f} rev/day < 11.25: not LEO. OrbitFM was trained on LEO only.") | |
| horizon = int(max(1, min(30, horizon_days))) | |
| n_ctx = int(max(32, min(256, context_days))) | |
| full = build_series(recs) | |
| first_day = full["grid"][0] | |
| last_day = full["grid"][-1] # noqa: F841 | |
| if anchor_date and anchor_date.strip(): | |
| try: | |
| anchor_u = datetime.strptime(anchor_date.strip(), "%Y-%m-%d").replace(tzinfo=timezone.utc).timestamp() | |
| except ValueError as e: | |
| raise gr.Error("Anchor date must be YYYY-MM-DD.") from e | |
| else: | |
| anchor_u = math.floor(recs[-1].epoch_unix / SECONDS_PER_DAY) * SECONDS_PER_DAY - horizon * SECONDS_PER_DAY | |
| # Anchor on an interior grid day of the full daily grid, as in the authors' eval | |
| # (grid days past the last TLE would have clamped/extrapolated phase). | |
| last_valid = int(np.searchsorted(full["grid"], recs[-1].epoch_unix, side="right") - 1) | |
| a = int(round((anchor_u - first_day) / SECONDS_PER_DAY)) | |
| if a > last_valid: | |
| a = last_valid | |
| if a < 0: | |
| raise gr.Error(f"Anchor date is before the data. Data covers {_date(first_day)} to {_date(recs[-1].epoch_unix)}.") | |
| ctx = full | |
| t_a = float(ctx["grid"][a]) | |
| if not ctx["mask"][a]: | |
| gr.Warning("No TLE within 2 days of the anchor date; the anchor state is interpolated across a gap.") | |
| if a + 1 < n_ctx: | |
| raise gr.Error(f"Need {n_ctx} days of history before the anchor; only {a + 1} available. " | |
| f"Data covers {_date(first_day)} to {_date(recs[-1].epoch_unix)}.") | |
| L = math.ceil(n_ctx / PATCH) * PATCH | |
| target = np.zeros((N_CHANNELS, L), dtype=np.float32) | |
| mask = np.zeros((N_CHANNELS, L), dtype=bool) | |
| window = ctx["feats"][a - n_ctx + 1: a + 1] | |
| cov = ctx["mask"][a - n_ctx + 1: a + 1] | |
| target[:, L - n_ctx:] = window.T | |
| mask[:N_ORBITAL, L - n_ctx:] = cov[None, :] | |
| mask[N_ORBITAL:, L - n_ctx:] = True | |
| coverage = float(cov.mean()) | |
| t_gpu = time.perf_counter() | |
| med, q10, q90 = _forecast_gpu(target, mask, horizon) | |
| gpu_s = time.perf_counter() - t_gpu | |
| anchor_feat, anchor_aux = ctx["feats"][a], ctx["aux"][a] | |
| anchor_full = elements_from_feat_aux(anchor_feat, anchor_aux) | |
| track_m = reconstruct_track(anchor_aux, med[:N_ORBITAL].T) | |
| base_orb = 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_orb, (horizon, 1))) | |
| base_sat = build_satrec(anchor_full, t_a) | |
| # truth lookup on the full (all-TLE) grid | |
| g_index = {int(round(g / SECONDS_PER_DAY)): i for i, g in enumerate(full["grid"])} | |
| rows = [] | |
| for h in range(1, horizon + 1): | |
| t_j = t_a + h * SECONDS_PER_DAY | |
| mm_eff = float(np.mean([track_m[k]["mean_motion_rev_per_day"] for k in range(h)])) | |
| r_m = propagate(build_satrec({**anchor_full, "mean_motion_rev_per_day": mm_eff, "bstar": 0.0}, t_a), t_j) | |
| r_b = propagate(base_sat, t_j) | |
| j = g_index.get(int(round(t_j / SECONDS_PER_DAY))) | |
| true_el, r_t = None, None | |
| if j is not None and j <= last_valid and full["mask"][j]: | |
| true_el = elements_from_feat_aux(full["feats"][j], full["aux"][j]) | |
| r_t = propagate(build_satrec(true_el, t_j), t_j) | |
| pm, pb = track_m[h - 1], track_b[h - 1] | |
| rows.append({ | |
| "day": h, "date": _date(t_j), | |
| "orbitfm_alt_km": float(mm_to_alt(pm["mean_motion_rev_per_day"])), | |
| "orbitfm_mean_motion": pm["mean_motion_rev_per_day"], | |
| "orbitfm_raan_deg": pm["raan_deg"], "orbitfm_argp_deg": pm["argp_deg"], | |
| "orbitfm_incl_deg": pm["inclination_deg"], "orbitfm_ecc": pm["eccentricity"], | |
| "true_alt_km": float(mm_to_alt(true_el["mean_motion_rev_per_day"])) if true_el else np.nan, | |
| "true_raan_deg": true_el["raan_deg"] if true_el else np.nan, | |
| "true_incl_deg": true_el["inclination_deg"] if true_el else np.nan, | |
| "orbitfm_pos_err_km": float(np.linalg.norm(r_m - r_t)) if (r_t is not None and r_m is not None) else np.nan, | |
| "sgp4_pos_err_km": float(np.linalg.norm(r_b - r_t)) if (r_t is not None and r_b is not None) else np.nan, | |
| "orbitfm_vs_sgp4_km": float(np.linalg.norm(r_m - r_b)) if (r_m is not None and r_b is not None) else np.nan, | |
| "orbitfm_x_km": r_m[0] if r_m is not None else np.nan, | |
| "orbitfm_y_km": r_m[1] if r_m is not None else np.nan, | |
| "orbitfm_z_km": r_m[2] if r_m is not None else np.nan, | |
| }) | |
| df = pd.DataFrame(rows) | |
| has_truth = df["sgp4_pos_err_km"].notna().any() | |
| # ---- figure 1: position error ---- | |
| fig1 = go.Figure() | |
| if has_truth: | |
| fig1.add_trace(go.Scatter(x=df["day"], y=df["orbitfm_pos_err_km"], mode="lines+markers", | |
| name="OrbitFM", line=dict(color="#e8590c", width=3))) | |
| fig1.add_trace(go.Scatter(x=df["day"], y=df["sgp4_pos_err_km"], mode="lines+markers", | |
| name="SGP4 persistence (last TLE)", line=dict(color="#1c7ed6", width=3, dash="dash"))) | |
| fig1.update_layout(title="Position error vs. truth TLE (TEME, km)", yaxis_title="km") | |
| else: | |
| fig1.add_trace(go.Scatter(x=df["day"], y=df["orbitfm_vs_sgp4_km"], mode="lines+markers", | |
| name="|OrbitFM - SGP4 persistence|", line=dict(color="#e8590c", width=3))) | |
| fig1.update_layout(title="No future TLEs to compare against: divergence of OrbitFM from SGP4 persistence (km)", | |
| yaxis_title="km") | |
| fig1.update_layout(xaxis_title="forecast day", template="plotly_white", height=380, | |
| legend=dict(orientation="h", y=-0.25), margin=dict(l=50, r=20, t=60, b=40)) | |
| # ---- figure 2: element panels ---- | |
| hist_lo = max(0, a - n_ctx + 1) | |
| h_idx = np.arange(hist_lo, a + 1) | |
| h_obs = ctx["mask"][h_idx] | |
| h_dates = [_dt(g) for g in ctx["grid"][h_idx][h_obs]] | |
| f_dates = [_dt(t_a + h * SECONDS_PER_DAY) for h in range(1, horizon + 1)] | |
| fig2 = make_subplots(rows=2, cols=2, subplot_titles=( | |
| "Mean altitude from mean motion (km)", "RAAN (deg, unwrapped from anchor)", | |
| "Daily mean-motion drift (rev/day²) · 10–90% band", "Inclination (deg)"), | |
| vertical_spacing=0.16, horizontal_spacing=0.09) | |
| # altitude | |
| fig2.add_trace(go.Scatter(x=h_dates, y=mm_to_alt(ctx["aux"][h_idx][h_obs, AUX_MM]), name="history (context)", | |
| line=dict(color="#868e96")), 1, 1) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=df["orbitfm_alt_km"], name="OrbitFM forecast", | |
| line=dict(color="#e8590c", width=3)), 1, 1) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=[mm_to_alt(anchor_full["mean_motion_rev_per_day"])] * horizon, | |
| name="SGP4 persistence", line=dict(color="#1c7ed6", dash="dash")), 1, 1) | |
| if has_truth: | |
| fig2.add_trace(go.Scatter(x=f_dates, y=df["true_alt_km"], name="truth (later TLEs)", | |
| line=dict(color="#2b8a3e", width=2), mode="lines+markers", | |
| marker=dict(size=4)), 1, 1) | |
| # RAAN unwrapped relative to anchor | |
| raan_hist = np.degrees(np.unwrap(np.radians(ctx["aux"][h_idx][h_obs, AUX_RAAN]))) | |
| raan0 = float(anchor_aux[AUX_RAAN]) | |
| raan_hist = raan_hist - raan_hist[-1] + raan0 if len(raan_hist) else raan_hist | |
| raan_f = raan0 + np.cumsum(med[4]) | |
| fig2.add_trace(go.Scatter(x=h_dates, y=raan_hist, showlegend=False, line=dict(color="#868e96")), 1, 2) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=raan_f, showlegend=False, line=dict(color="#e8590c", width=3)), 1, 2) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=[raan0] * horizon, showlegend=False, | |
| line=dict(color="#1c7ed6", dash="dash")), 1, 2) | |
| if has_truth: | |
| tr = df["true_raan_deg"].to_numpy() | |
| ok = np.isfinite(tr) | |
| tr_uw = np.full_like(tr, np.nan) | |
| if ok.any(): | |
| tr_uw[ok] = raan0 + ((tr[ok] - raan0 + 180.0) % 360.0 - 180.0) | |
| tr_uw[ok] = np.degrees(np.unwrap(np.radians(tr_uw[ok]))) | |
| # align branch with model forecast | |
| tr_uw[ok] += 360.0 * np.round((raan_f[ok] - tr_uw[ok]) / 360.0) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=tr_uw, showlegend=False, mode="lines+markers", | |
| marker=dict(size=4), line=dict(color="#2b8a3e", width=2)), 1, 2) | |
| # d mean motion with band | |
| fig2.add_trace(go.Scatter(x=h_dates, y=ctx["feats"][h_idx][h_obs, 1], showlegend=False, line=dict(color="#868e96")), 2, 1) | |
| fig2.add_trace(go.Scatter(x=f_dates + f_dates[::-1], y=list(q90[1]) + list(q10[1][::-1]), fill="toself", | |
| fillcolor="rgba(232,89,12,0.18)", line=dict(width=0), showlegend=False, | |
| hoverinfo="skip"), 2, 1) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=med[1], showlegend=False, line=dict(color="#e8590c", width=3)), 2, 1) | |
| # inclination | |
| fig2.add_trace(go.Scatter(x=h_dates, y=ctx["feats"][h_idx][h_obs, 3], showlegend=False, | |
| line=dict(color="#868e96")), 2, 2) | |
| fig2.add_trace(go.Scatter(x=f_dates, y=df["orbitfm_incl_deg"], showlegend=False, | |
| line=dict(color="#e8590c", width=3)), 2, 2) | |
| if has_truth: | |
| ti = df["true_incl_deg"].to_numpy() | |
| fig2.add_trace(go.Scatter(x=f_dates, y=ti, showlegend=False, mode="lines+markers", | |
| marker=dict(size=4), line=dict(color="#2b8a3e", width=2)), 2, 2) | |
| for rr, cc in [(1, 1), (1, 2), (2, 1), (2, 2)]: | |
| fig2.add_vline(x=_dt(t_a).timestamp() * 1000, line=dict(color="#adb5bd", dash="dot"), row=rr, col=cc) | |
| fig2.update_layout(template="plotly_white", height=640, legend=dict(orientation="h", y=-0.08), | |
| margin=dict(l=50, r=20, t=50, b=40)) | |
| # ---- summary ---- | |
| title = name or f"NORAD {norad}" | |
| lines = [f"### {title} (NORAD {norad})", | |
| f"Anchor **{_date(t_a)}** · context {n_ctx} d (coverage {coverage:.0%}) · horizon {horizon} d · " | |
| f"{len(recs)} TLEs used ({n_removed} outliers removed) · median altitude ≈ {float(mm_to_alt(med_mm)):.0f} km"] | |
| if has_truth: | |
| lines.append("\n| horizon | OrbitFM error (km) | SGP4 persistence error (km) |\n|---|---|---|") | |
| for h in [1, 3, 7, 14, 30]: | |
| if h <= horizon: | |
| r = df.iloc[h - 1] | |
| if np.isfinite(r["sgp4_pos_err_km"]): | |
| win = " ✅" if r["orbitfm_pos_err_km"] < r["sgp4_pos_err_km"] else "" | |
| lines.append(f"| {h} d | {r['orbitfm_pos_err_km']:.1f}{win} | {r['sgp4_pos_err_km']:.1f} |") | |
| else: | |
| lines.append(f"\nPure forecast (no TLEs after the anchor). Day-{horizon} OrbitFM vs SGP4 persistence " | |
| f"separation: **{df['orbitfm_vs_sgp4_km'].iloc[-1]:.1f} km**.") | |
| lines.append(f"\n<sub>Model forward pass {gpu_s:.2f}s · total {time.perf_counter() - t_start:.1f}s · " | |
| f"space weather: {SW_SRC}</sub>") | |
| summary = "\n".join(lines) | |
| table = df[["day", "date", "orbitfm_alt_km", "orbitfm_mean_motion", "orbitfm_raan_deg", | |
| "orbitfm_incl_deg", "orbitfm_ecc", "orbitfm_pos_err_km", "sgp4_pos_err_km", | |
| "orbitfm_vs_sgp4_km"]].round(6) | |
| csv_path = f"/tmp/orbitfm_{norad}_{_date(t_a)}_{horizon}d.csv" | |
| df.to_csv(csv_path, index=False) | |
| return summary, fig1, fig2, table, csv_path | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# 🛰️ OrbitFM: satellite orbit forecasting from TLE history\n" | |
| "[OrbitFM](https://huggingface.co/postcn/OrbitFM) is a fine-tune of " | |
| "[Toto-2.0-2.5B](https://huggingface.co/Datadog/Toto-2.0-2.5B). It reads a LEO satellite's TLE history " | |
| "plus space weather (F10.7, Ap) and forecasts how its orbital elements will drift. Positions come from " | |
| "SGP4, and are compared with the standard *SGP4 persistence* baseline, which propagates the last TLE. " | |
| "Upload or paste a TLE history (e.g. from Space-Track), or try one of the examples below." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| tle_file = gr.File(label="TLE history (.txt / .tle, 2- or 3-line format)", | |
| file_types=[".txt", ".tle", ".3le"], type="filepath") | |
| anchor = gr.Textbox(label="Anchor / forecast start date (YYYY-MM-DD, UTC)", | |
| placeholder="empty = backtest ending at the last TLE", value="") | |
| horizon = gr.Slider(1, 30, value=30, step=1, label="Forecast horizon (days)") | |
| with gr.Accordion("Advanced", open=False): | |
| context = gr.Slider(32, 256, value=64, step=1, | |
| label="Context length (days). 64 matches the paper's evaluation") | |
| tle_text = gr.Textbox(label="…or paste TLE text (overrides the file)", lines=6, | |
| placeholder="1 20580U 90037B ...\n2 20580 28.4665 ...") | |
| run = gr.Button("Forecast orbit", variant="primary") | |
| with gr.Column(scale=2): | |
| summary = gr.Markdown() | |
| fig_err = gr.Plot(label="Position error") | |
| fig_el = gr.Plot(label="Orbital elements: history, forecast, truth") | |
| with gr.Accordion("Per-day forecast table", open=False): | |
| table = gr.Dataframe(wrap=False) | |
| csv_out = gr.File(label="Download full forecast (CSV)") | |
| outputs = [summary, fig_err, fig_el, table, csv_out] | |
| run.click(forecast_orbit, inputs=[tle_file, anchor, horizon, context, tle_text], | |
| outputs=outputs, api_name="forecast") | |
| gr.Examples( | |
| examples=[ | |
| ["examples/hubble_20580.txt", "2026-09-01", 30], | |
| ["examples/pegasus_rb_27371.txt", "2026-08-20", 30], | |
| ["examples/envisat_27386.txt", "2026-09-05", 30], | |
| ["examples/noaa19_33591.txt", "2026-09-01", 14], | |
| ["examples/hubble_20580.txt", "2026-10-09", 30], | |
| ], | |
| inputs=[tle_file, anchor, horizon], | |
| outputs=outputs, | |
| fn=forecast_orbit, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Real Space-Track TLE histories (Jun 2025 – Oct 2026, after OrbitFM's training data)", | |
| ) | |
| gr.Markdown( | |
| "<sub>Example TLEs: [juliensimon/space-track-tle-history](https://huggingface.co/datasets/juliensimon/" | |
| "space-track-tle-history) (CC-BY-4.0, from Space-Track.org / 18th SDS). Space weather: CelesTrak SW-All. " | |
| "This is a research model; don't use it for collision avoidance or other operational decisions.</sub>" | |
| ) | |
| demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |