#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ TLE daily-grid dataset v2: outlier-cleaned + space-weather channels. Builds on the v1 daily-grid cumulative-phase formulation and adds the two data-quality / physics levers from the cm-tle-pred benchmark analysis: (1) Robust outlier removal per satellite (tle_clean.clean_records) before resampling -- their single biggest accuracy lever. (2) Space-weather input channels (F10.7, F10.7-81d, Ap) -- exogenous drivers of atmospheric drag, i.e. the secular decay of mean motion. Channels (CHANNEL_NAMES): orbital, PREDICTED + in loss (indices 0..5): bstar_slog, mean_motion, eccentricity, inclination_deg, draan_deg_per_day, dargp_deg_per_day solar, INPUT-ONLY context, excluded from loss (indices 6..8): f107, f107_81, ap Absolute angles [MA, RAAN, argp] are kept as aux (anchor/truth), never predicted. Angles are reconstructed by anchoring at the last observed TLE and integrating the predicted rates (reconstruct_track) -- absolute phase is never predicted. """ from __future__ import annotations import hashlib import json import math import sys from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Iterable, List, Optional import numpy as np import torch from torch.utils.data import Dataset from tqdm import tqdm _HERE = Path(__file__).resolve().parent _CODE_UTILS = _HERE.parents[1] / "code" / "utils" for p in (str(_HERE), str(_CODE_UTILS)): if p not in sys.path: sys.path.insert(0, p) from tle_future_dataset import ( # noqa: E402 collect_txt_files, iter_tle_pairs_from_txt, parse_tle_pair, parse_date_to_unix, signed_log1p, extract_year_from_filename, ) from tle_clean import clean_records, cumulative_mean_anomaly # noqa: E402 from space_weather import SpaceWeather, load_space_weather, SOLAR_FEATURES, N_SOLAR # noqa: E402 # v3: persistence-RESIDUAL (drift) targets. # # v2 predicted absolute mean_motion / bstar. Because these barely change over the # forecast window, the loss-optimal prediction was ≈persistence and the model # could not beat SGP4 propagation at short/mid horizons. v3 makes the slowly # DRIFTING magnitudes into per-day DELTA channels, so the training target IS the # drift (zero-mean, small) — the learning signal we care about is amplified. # Absolute mean_motion / bstar at the anchor live in aux, and the trajectory is # reconstructed by anchor + cumulative predicted delta (like the angle rates). # # Eccentricity / inclination drift even less and are well predicted as absolutes, # so they stay absolute. RAAN/argp remain secular rate channels (as in v2). ORBITAL_NAMES = [ "d_bstar_slog_per_day", # 0 daily drift of signed_log1p(bstar) "d_mean_motion_per_day", # 1 daily drift of mean motion (orbital decay rate) "eccentricity", # 2 absolute "inclination_deg", # 3 absolute "draan_deg_per_day", # 4 RAAN secular rate "dargp_deg_per_day", # 5 argp secular rate ] N_ORBITAL = len(ORBITAL_NAMES) # v4 = cm-tle-pred-style processing. The cm-tle-pred benchmark feeds, on top of the # raw elements, a set of engineered physical features. We add the reproducible ones # as INPUT-ONLY channels (excluded from the loss), mirroring their feature names: # - SAT_RX..VZ : Cartesian ECI state from the day's elements (two-body # osculating r,v -- same role as their SGP4 SAT_R*/V*) # - SEMIMAJOR_AXIS / PERIOD / APOAPSIS / PERIAPSIS : derived orbit geometry # - *_COS / *_SIN : cyclical encodings of the angular elements (MA/RAAN/argp) # What is PREDICTED is unchanged (the 6 orbital drift/rate channels) -- this only # enriches Toto's context, exactly like the solar channels do. PHYS_NAMES = [ "sat_rx", "sat_ry", "sat_rz", # 6..8 ECI position (km) "sat_vx", "sat_vy", "sat_vz", # 9..11 ECI velocity (km/s) "semimajor_axis", "period_min", # 12..13 (km), (minutes/rev) "apoapsis_alt", "periapsis_alt", # 14..15 altitude (km) "ma_cos", "ma_sin", # 16..17 mean anomaly cyclical "raan_cos", "raan_sin", # 18..19 RAAN cyclical "argp_cos", "argp_sin", # 20..21 argp cyclical ] N_PHYS = len(PHYS_NAMES) CHANNEL_NAMES = ORBITAL_NAMES + PHYS_NAMES + SOLAR_FEATURES # 6 + 16 + 3 = 25 N_CHANNELS = len(CHANNEL_NAMES) # aux carries the absolute anchors/truth NOT directly predicted: # [mean_anomaly, RAAN, argp, mean_motion, bstar_slog] N_AUX = 5 AUX_MA, AUX_RAAN, AUX_ARGP, AUX_MM, AUX_BSTAR = 0, 1, 2, 3, 4 # orbital channels are prediction targets; PHYS + solar are input-only context. LOSS_CHANNEL_MASK = np.array([1] * N_ORBITAL + [0] * (N_PHYS + N_SOLAR), dtype=bool) # Position-critical drift channels: d_bstar_slog (0) and d_mean_motion (1). SGP4 # along-track error is governed by mean motion and drag, so these are upweighted in # the training loss (--drift-loss-weight) to align the objective with position. DRIFT_CHANNELS = (0, 1) SECONDS_PER_DAY = 86400.0 MU_EARTH = 398600.4418 # km^3/s^2 (WGS-72 gravitational parameter) R_EARTH = 6378.137 # km (equatorial radius, for apo/peri altitude) LEO_MIN_MEAN_MOTION = 11.25 # rev/day; cm-tle-pred LEO cutoff CHANNEL_VERSION = "v5cmtle_feats" def signed_expm1(x: float) -> float: return math.copysign(math.expm1(abs(x)), x) def _clip_ecc(x): return float(min(max(x, 0.0), 0.9999999)) def _clip_inc(x): return float(min(max(x, 0.0), 180.0)) def elements_from_feat_aux(feat_row, aux_row) -> Dict[str, float]: """Full physical elements for an OBSERVED grid row (truth/anchor). ecc/inc come from the absolute channels; bstar/mean_motion/angles from aux. """ f = np.asarray(feat_row, dtype=np.float64) a = np.asarray(aux_row, dtype=np.float64) return { "bstar": signed_expm1(float(a[AUX_BSTAR])), "mean_motion_rev_per_day": float(max(a[AUX_MM], 1e-6)), "eccentricity": _clip_ecc(f[2]), "inclination_deg": _clip_inc(f[3]), "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_seq, grid_step_days=1.0): """Absolute elements per forecast day, anchored at the last observed TLE. bstar(log) and mean_motion are integrated from their anchor (aux) by summing the predicted per-day DELTAS; mean-anomaly phase is integrated trapezoidally from the (reconstructed) mean motion; RAAN/argp from predicted daily rates; ecc/inc are read from the absolute channels. Nothing absolute-phase is predicted. """ a = np.asarray(anchor_aux, dtype=np.float64) pred = np.asarray(pred_seq, dtype=np.float64) ma, raan, argp = a[AUX_MA], a[AUX_RAAN], a[AUX_ARGP] mm = float(a[AUX_MM]); bstar_slog = float(a[AUX_BSTAR]) prev_mm = mm out = [] for h in range(pred.shape[0]): bstar_slog += float(pred[h, 0]) * grid_step_days # cumulative d_bstar mm += float(pred[h, 1]) * grid_step_days # cumulative d_mean_motion mm_h = max(mm, 1e-6) ma += 360.0 * 0.5 * (prev_mm + mm_h) * grid_step_days prev_mm = mm_h raan += float(pred[h, 4]) * grid_step_days argp += float(pred[h, 5]) * grid_step_days out.append({ "bstar": signed_expm1(bstar_slog), "mean_motion_rev_per_day": mm_h, "eccentricity": _clip_ecc(pred[h, 2]), "inclination_deg": _clip_inc(pred[h, 3]), "mean_anomaly_deg": ma % 360.0, "raan_deg": raan % 360.0, "argp_deg": argp % 360.0, }) return out # ----------------------------- # cm-tle-pred physical features (input-only channels) # ----------------------------- def _kepler_rv(a, e, inc_r, raan_r, argp_r, M_r): """Vectorized two-body osculating ECI state from classical elements. Inputs are numpy arrays (one element per grid day); returns rx,ry,rz (km) and vx,vy,vz (km/s). This plays the same role as cm-tle-pred's SGP4 SAT_R*/V* Cartesian features but is fully vectorized over the daily grid (no per-day sgp4init), which matters for the all-years (~23 GB) build. """ E = np.array(M_r, dtype=np.float64, copy=True) for _ in range(12): # Newton on Kepler's equation 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) # perifocal -> ECI rotation R3(-raan) R1(-inc) R3(-argp) 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 rx = R11 * rp_x + R12 * rp_y ry = R21 * rp_x + R22 * rp_y rz = R31 * rp_x + R32 * rp_y vx = R11 * vp_x + R12 * vp_y vy = R21 * vp_x + R22 * vp_y vz = R31 * vp_x + R32 * vp_y return rx, ry, rz, vx, vy, vz def physics_features(g_mm, g_ecc, g_inc_deg, g_raan_deg, g_argp_deg, g_ma_deg): """(T, N_PHYS) cm-tle-pred input features from the daily-grid elements.""" mm = np.maximum(np.asarray(g_mm, dtype=np.float64), 1e-6) n_rad_s = mm * 2.0 * math.pi / SECONDS_PER_DAY a = (MU_EARTH / (n_rad_s ** 2)) ** (1.0 / 3.0) # semimajor axis (km) e = np.clip(np.asarray(g_ecc, dtype=np.float64), 0.0, 0.999999) inc_r = np.radians(g_inc_deg) raan_r = np.radians(np.asarray(g_raan_deg) % 360.0) argp_r = np.radians(np.asarray(g_argp_deg) % 360.0) ma_r = np.radians(np.asarray(g_ma_deg) % 360.0) rx, ry, rz, vx, vy, vz = _kepler_rv(a, e, inc_r, raan_r, argp_r, ma_r) period_min = 1440.0 / mm apo = a * (1.0 + e) - R_EARTH # apoapsis altitude (km) peri = a * (1.0 - e) - R_EARTH # periapsis altitude (km) return np.stack([ rx, ry, rz, vx, vy, vz, a, period_min, apo, peri, 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 sat_split_of(norad_id, train_frac=0.70, valid_frac=0.15): """cm-tle-pred-style SATELLITE-level split (70/15/15): every record of a given satellite lands in the same split, deterministically by NORAD id.""" h = int(hashlib.md5(str(int(norad_id)).encode()).hexdigest(), 16) % 1000 if h < int(train_frac * 1000): return "train" if h < int((train_frac + valid_frac) * 1000): return "valid" return "test" # ----------------------------- # Parse + clean (with progress) # ----------------------------- def _is_physically_possible(r) -> bool: """cm-tle-pred 'data integrity' check: drop physically impossible records.""" return (0.0 <= r.eccentricity < 1.0 and 0.0 < r.inclination_deg < 180.0 and 0.1 < r.mean_motion_rev_per_day < 20.0) def load_and_clean( input_dir, years: Optional[Iterable[int]] = None, clean: bool = True, min_records_per_object: int = 2, drop_first_n: int = 5, leo_only: bool = True, ) -> Dict[int, List]: """Parse raw TLE txt (tqdm over files), group by NORAD, dedup, outlier-clean. v4 adds the cm-tle-pred 'data integrity' steps: - drop physically impossible records (impossible ecc/inc/mean_motion), - discard the first ``drop_first_n`` TLEs per satellite (least reliable), - keep only LEO objects (median mean_motion >= LEO_MIN_MEAN_MOTION) when ``leo_only`` is set. """ files = collect_txt_files(input_dir, years=years) by_norad: Dict[int, List] = {} for path in tqdm(files, desc="parse files", unit="file"): year = extract_year_from_filename(path) ridx = 0 for obj, l1, l2, ln in iter_tle_pairs_from_txt(path): ridx += 1 rec = parse_tle_pair(l1, l2, source_file=str(path), source_year=year, source_line_number=ln, object_name_raw=obj, record_index=ridx) if rec is not None and _is_physically_possible(rec): by_norad.setdefault(rec.norad_id, []).append(rec) out: Dict[int, List] = {} removed_total = 0 n_nonleo = 0 for norad, recs in tqdm(by_norad.items(), desc="dedup+clean", unit="sat"): recs = sorted(recs, key=lambda r: r.epoch_unix) deduped, seen = [], set() for r in recs: key = round(r.epoch_unix, 3) if key in seen: continue seen.add(key) deduped.append(r) if drop_first_n > 0 and len(deduped) > drop_first_n: # least-reliable early TLEs deduped = deduped[drop_first_n:] if leo_only: med_mm = float(np.median([r.mean_motion_rev_per_day for r in deduped])) if deduped else 0.0 if med_mm < LEO_MIN_MEAN_MOTION: n_nonleo += 1 continue if clean and len(deduped) >= 5: deduped, nrem = clean_records(deduped) removed_total += nrem if len(deduped) >= min_records_per_object: out[norad] = deduped print(f"[load] {len(out)} satellites, removed {removed_total} outlier records, " f"dropped {n_nonleo} non-LEO objects (leo_only={leo_only})") return out # ----------------------------- # Daily-grid series + solar # ----------------------------- @dataclass class DailySeries: norad_id: int grid_epochs: np.ndarray # (T,) feats: np.ndarray # (T, N_CHANNELS) mask: np.ndarray # (T,) orbital coverage aux: np.ndarray # (T, N_AUX) absolute angles def _build_one(recs, sw: Optional[SpaceWeather], grid_step_days=1.0) -> Optional[DailySeries]: 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], dtype=np.float64) 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 = grid_step_days * SECONDS_PER_DAY t0 = math.floor(ep[0] / step) * step t1 = math.ceil(ep[-1] / step) * step grid = np.arange(t0, t1 + 0.5 * step, step, dtype=np.float64) if grid.size < 2: return None def itp(y): return np.interp(grid, ep, y) g_raan, g_argp = itp(raan_uw), itp(argp_uw) g_mm, g_bstar = itp(mm), itp(bstar_log) g_ecc, g_inc, g_ma = itp(ecc), itp(inc), itp(phiM) # g_ma = cumulative MA (deg) draan = np.gradient(g_raan) / grid_step_days dargp = np.gradient(g_argp) / grid_step_days d_bstar = np.gradient(g_bstar) / grid_step_days # daily drift (persistence residual) d_mm = np.gradient(g_mm) / grid_step_days orbital = np.stack([d_bstar, d_mm, g_ecc, g_inc, draan, dargp], axis=1) # cm-tle-pred input-only physical features (Cartesian r/v, SMA/period/apo/peri, # cyclical angle encodings) derived from the day's elements phys = physics_features(g_mm, g_ecc, g_inc, g_raan, g_argp, g_ma) if sw is not None: solar = sw.for_epochs(grid) # (T, N_SOLAR) else: solar = np.zeros((grid.size, N_SOLAR), dtype=np.float32) feats = np.concatenate([orbital, phys, solar], axis=1).astype(np.float32) # aux = [MA, RAAN, argp, mean_motion_abs, bstar_slog_abs] 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])) mask = gap <= (2.0 * step) return DailySeries(recs[0].norad_id, grid, feats, mask, aux) def _cache_key(input_dir, years, grid_step_days, clean, with_solar, min_grid_points, leo_only) -> str: h = hashlib.md5(str(Path(input_dir).resolve()).encode()).hexdigest()[:6] tag = Path(input_dir).resolve().name ys = "all" if years is None else f"{min(years)}-{max(years)}_{len(list(years))}" flags = f"c{int(clean)}s{int(with_solar)}m{int(min_grid_points)}L{int(leo_only)}" return f"{tag}_{h}_{ys}_g{grid_step_days:g}_{CHANNEL_VERSION}_{flags}" def _load_cache_npz(cache_path, verbose) -> Dict[int, DailySeries]: d = np.load(cache_path, allow_pickle=True) out = {int(n): DailySeries(int(n), d[f"e_{int(n)}"], d[f"f_{int(n)}"], d[f"m_{int(n)}"], d[f"a_{int(n)}"]) for n in d["norads"]} if verbose: print(f"[cache] {len(out)} satellites loaded from {cache_path}") return out def build_daily_series( input_dir, years=None, grid_step_days=1.0, min_grid_points=64, clean=True, sw_csv=None, cache_dir=None, cache_file=None, rebuild=False, leo_only=True, verbose=True, ) -> Dict[int, DailySeries]: # Explicit prebuilt cache (e.g. the full 2005-2024 superset): load it directly # and let the dataset/eval filter by window + time-split. Skips parsing & solar. if cache_file is not None: cp = Path(cache_file) if not cp.exists(): raise FileNotFoundError(f"--cache-file not found: {cp}") if verbose: print(f"[cache] using explicit cache {cp}") return _load_cache_npz(cp, verbose) sw = load_space_weather(sw_csv) with_solar = sw is not None cache_path = None if cache_dir is not None: cache_dir = Path(cache_dir); cache_dir.mkdir(parents=True, exist_ok=True) cache_path = cache_dir / f"tle_v4_{_cache_key(input_dir, years, grid_step_days, clean, with_solar, min_grid_points, leo_only)}.npz" if cache_path.exists() and not rebuild: if verbose: print(f"[cache] loading {cache_path}") d = np.load(cache_path, allow_pickle=True) out = {int(n): DailySeries(int(n), d[f"e_{int(n)}"], d[f"f_{int(n)}"], d[f"m_{int(n)}"], d[f"a_{int(n)}"]) for n in d["norads"]} if verbose: print(f"[cache] {len(out)} satellites loaded") return out by_norad = load_and_clean(input_dir, years=years, clean=clean, leo_only=leo_only) out: Dict[int, DailySeries] = {} for norad, recs in tqdm(by_norad.items(), desc="daily grid+solar", unit="sat"): ser = _build_one(recs, sw, grid_step_days) if ser is None or int(ser.mask.sum()) < min_grid_points: continue out[norad] = ser if cache_path is not None: kw: Dict[str, Any] = {"norads": np.asarray(sorted(out.keys()))} for n, s in out.items(): kw[f"e_{n}"] = s.grid_epochs; kw[f"f_{n}"] = s.feats kw[f"m_{n}"] = s.mask; kw[f"a_{n}"] = s.aux np.savez(cache_path, **kw) if verbose: print(f"[cache] saved {len(out)} satellites -> {cache_path}") return out # ----------------------------- # Windowed dataset # ----------------------------- @dataclass class DayWindow: norad_id: int start: int end_epoch: float class TLEDatasetV2(Dataset): def __init__( self, input_dir, patch_size, years=None, window_patches=4, stride_patches=2, min_valid_steps=None, min_grid_points=None, split="train", train_until="2022-01-01", valid_until="2023-01-01", clean=True, sw_csv=None, cache_dir=None, cache_file=None, rebuild_cache=False, max_satellites=None, grid_step_days=1.0, leo_only=True, split_mode="time", verbose=True, ): if split not in {"train", "valid", "test", "all"}: raise ValueError("split must be train|valid|test|all") if split_mode not in {"time", "satellite"}: raise ValueError("split_mode must be time|satellite") self.patch_size = int(patch_size) self.window_length = int(window_patches) * self.patch_size self.stride = max(1, int(stride_patches) * self.patch_size) self.min_obs = int(min_valid_steps or self.window_length) self.split = split self.split_mode = split_mode # 'time' (epoch cutoffs) or 'satellite' (cm-tle-pred 70/15/15) self.train_until_unix = parse_date_to_unix(train_until) self.valid_until_unix = parse_date_to_unix(valid_until) # cache-build satellite filter: keep a satellite if it has at least this many # observed grid days. Decoupled from window length so the cache can hold the # full population (per-window length is enforced later in _build_index). cache_min = int(min_grid_points) if min_grid_points is not None else self.window_length self.series = build_daily_series( input_dir, years=years, grid_step_days=grid_step_days, min_grid_points=cache_min, clean=clean, sw_csv=sw_csv, cache_dir=cache_dir, cache_file=cache_file, rebuild=rebuild_cache, leo_only=leo_only, verbose=verbose, ) if max_satellites is not None: keep = sorted(self.series.keys())[:max_satellites] self.series = {k: self.series[k] for k in keep} self.index: List[DayWindow] = self._build_index() if verbose: print(json.dumps({ "dataset": "TLEDatasetV2", "split": split, "split_mode": split_mode, "num_satellites": len(self.series), "num_windows": len(self.index), "window_days": self.window_length, "patch_size": self.patch_size, "n_channels": N_CHANNELS, "orbital": N_ORBITAL, "phys": N_PHYS, "solar": N_SOLAR, }, indent=2)) def _split_of(self, epoch): if self.train_until_unix is None or self.valid_until_unix is None: return "all" if epoch < self.train_until_unix: return "train" if epoch < self.valid_until_unix: return "valid" return "test" def _build_index(self): out, W = [], self.window_length for norad, s in self.series.items(): T = s.feats.shape[0] if T < W: continue # satellite-level split: the whole satellite belongs to one split if self.split != "all" and self.split_mode == "satellite" \ and sat_split_of(norad) != self.split: continue for st in range(0, T - W + 1, self.stride): if int(s.mask[st:st + W].sum()) < self.min_obs: continue end_epoch = float(s.grid_epochs[st + W - 1]) if self.split != "all" and self.split_mode == "time" \ and self._split_of(end_epoch) != self.split: continue out.append(DayWindow(norad, st, end_epoch)) return out def __len__(self): return len(self.index) def __getitem__(self, i): w = self.index[i] s = self.series[w.norad_id] W = self.window_length feats = s.feats[w.start:w.start + W] mask = s.mask[w.start:w.start + W] epochs = s.grid_epochs[w.start:w.start + W] aux = s.aux[w.start:w.start + W] # orbital channels use coverage mask; solar channels are always observed chan_mask = np.zeros((N_CHANNELS, W), dtype=bool) chan_mask[:N_ORBITAL] = mask[None, :] chan_mask[N_ORBITAL:] = True return { "target": torch.from_numpy(feats.T.copy()), # (C, W) "target_mask": torch.from_numpy(chan_mask), # (C, W) "series_ids": torch.zeros(N_CHANNELS, dtype=torch.long), "meta": {"norad_id": w.norad_id, "length": int(mask.sum()), "epochs": epochs.copy(), "feats": feats.copy(), "aux": aux.copy()}, } def series_collate_fn(batch): return { "target": torch.stack([b["target"] for b in batch], 0), "target_mask": torch.stack([b["target_mask"] for b in batch], 0), "series_ids": torch.stack([b["series_ids"] for b in batch], 0), "meta": [b["meta"] for b in batch], } def main(): import argparse from torch.utils.data import DataLoader ap = argparse.ArgumentParser() 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("--sw-csv", default="/home/irteam/data-vol1/models/OrbitGPT/v2/data/SW-All.csv") ap.add_argument("--start-year", type=int, default=2020) ap.add_argument("--end-year", type=int, default=2020) ap.add_argument("--patch-size", type=int, default=32) ap.add_argument("--window-patches", type=int, default=3) ap.add_argument("--stride-patches", type=int, default=1) ap.add_argument("--min-grid-points", type=int, default=64, help="cache-build filter: keep satellites with >= this many observed " "grid days. Low value (e.g. 64) keeps the full population.") ap.add_argument("--no-clean", action="store_true") ap.add_argument("--no-leo", action="store_true", help="disable LEO-only filter") ap.add_argument("--split", default="all") ap.add_argument("--split-mode", default="time", choices=["time", "satellite"]) args = ap.parse_args() ds = TLEDatasetV2( input_dir=args.input_dir, cache_dir=args.cache_dir, sw_csv=args.sw_csv, years=range(args.start_year, args.end_year + 1), patch_size=args.patch_size, window_patches=args.window_patches, stride_patches=args.stride_patches, min_grid_points=args.min_grid_points, clean=not args.no_clean, leo_only=not args.no_leo, split=args.split, split_mode=args.split_mode, ) if len(ds) == 0: print("No windows."); return b = next(iter(DataLoader(ds, batch_size=4, shuffle=True, collate_fn=series_collate_fn))) print("target", tuple(b["target"].shape), "obs frac", round(float(b["target_mask"].float().mean()), 3)) print("channels:", N_CHANNELS, "= orbital", N_ORBITAL, "+ phys", N_PHYS, "+ solar", N_SOLAR) m0 = b["meta"][0] print("phys[t0] (rx,ry,rz,vx,vy,vz,...):", [round(float(x), 3) for x in m0["feats"][0][N_ORBITAL:N_ORBITAL + N_PHYS]]) print("solar[t0] (f107,f107_81,ap):", m0["feats"][0][N_ORBITAL + N_PHYS:].tolist()) track = reconstruct_track(m0["aux"][0], m0["feats"][1:6]) print("reconstructed[+5d]:", json.dumps(track[-1], indent=2)) print("true[+5d]:", json.dumps(elements_from_feat_aux(m0["feats"][5], m0["aux"][5]), indent=2)) if __name__ == "__main__": main()