scripts/eval_storm_conditioned.py

#!/usr/bin/env python3
"""Assess models under different geomagnetic conditions.

Usage:
    python scripts/eval_storm_conditioned.py --spacecraft iss
    python scripts/eval_storm_conditioned.py --spacecraft all
"""
import argparse
import json
import logging
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import yaml

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from src.data.preprocessing import OrbitPreprocessor

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger("storm-assessment")

RESULTS_DIR = Path("results")
CHECKPOINT_DIR = Path("checkpoints")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_config():
    with open("config.yaml") as f:
        return yaml.safe_load(f)


def load_model(model_type, checkpoint_path, input_dim=6, solar_dim=8):
    """Load model, auto-detecting architecture from checkpoint weights."""
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    state = ckpt["model_state_dict"]

    from scripts.train_gpu import OrbitLSTMDirect, OrbitTransformerDirect, SolarWindOrbitModel

    if model_type == "lstm":
        hidden_dim = state["lstm.weight_ih_l0"].shape[0] // 4
        layer_keys = [k for k in state if k.startswith("lstm.weight_ih_l")]
        num_layers = len(layer_keys) // 2
        fc_keys = sorted([k for k in state if k.startswith("fc.") and k.endswith(".weight")])
        horizon = state[fc_keys[-1]].shape[0] // 3
        model = OrbitLSTMDirect(input_dim=input_dim, hidden_dim=hidden_dim,
                                num_layers=num_layers, horizon=horizon, output_dim=3, dropout=0.0)

    elif model_type == "transformer":
        hidden_dim = state["input_proj.weight"].shape[0]
        layer_keys = [k for k in state if "layers." in k and "self_attn.in_proj_weight" in k]
        num_layers = len(layer_keys)
        # nhead must divide hidden_dim; train_gpu.py uses 8
        nhead = 8 if hidden_dim % 8 == 0 else 4
        fc_keys = sorted([k for k in state if k.startswith("head.") and k.endswith(".weight")])
        horizon = state[fc_keys[-1]].shape[0] // 3
        ff_dim = state["encoder.layers.0.linear1.weight"].shape[0]
        model = OrbitTransformerDirect(input_dim=input_dim, d_model=hidden_dim, nhead=nhead,
                                       num_layers=num_layers, dim_feedforward=ff_dim,
                                       horizon=horizon, output_dim=3, dropout=0.0)

    elif model_type == "multimodal":
        hidden_dim = state["orbit_enc.weight_ih_l0"].shape[0] // 4
        layer_keys = [k for k in state if k.startswith("orbit_enc.weight_ih_l")]
        num_layers = len(layer_keys) // 2
        fc_keys = sorted([k for k in state if k.startswith("base_head.") and k.endswith(".weight")])
        horizon = state[fc_keys[-1]].shape[0] // 3
        model = SolarWindOrbitModel(orbit_input_dim=input_dim, solar_input_dim=solar_dim,
                                     hidden_dim=hidden_dim, num_layers=num_layers, nhead=8,
                                     horizon=horizon, output_dim=3, dropout=0.0)
    else:
        raise ValueError(f"Unknown model: {model_type}")

    model.load_state_dict(state)
    model.to(DEVICE)
    model.eval()
    return model


def assign_kp(window_times, solar_df):
    """Assign preceding Kp to each window start time."""
    kp_df = solar_df[["time", "kp"]].dropna(subset=["kp"]).copy()
    kp_df["time"] = pd.to_datetime(kp_df["time"], utc=True).dt.tz_localize(None).astype("datetime64[ns]")
    kp_df = kp_df.sort_values("time").drop_duplicates("time")
    windows_df = pd.DataFrame({"time": pd.to_datetime(window_times).astype("datetime64[ns]")})
    windows_df["time"] = windows_df["time"].dt.tz_localize(None)
    windows_df = windows_df.sort_values("time")
    merged = pd.merge_asof(windows_df, kp_df, on="time", direction="backward")
    kp_vals = merged["kp"].values
    if len(kp_vals) == 0:
        return kp_vals
    # OMNI stores Kp*10 (0-90 scale). Convert to standard 0-9 scale.
    if np.nanmax(kp_vals[np.isfinite(kp_vals)]) > 9:
        kp_vals = kp_vals / 10.0
    return kp_vals


def run_spacecraft(spacecraft_id, config):
    log.info(f"=== {spacecraft_id} ===")

    proc = OrbitPreprocessor()
    df = pd.read_parquet(f"data/raw/{spacecraft_id}_2023-01-01_2025-12-31.parquet")
    processed = proc.preprocess(df, spacecraft_id)
    stats = proc.stats

    stats_dir = RESULTS_DIR / "norm_stats"
    stats_dir.mkdir(parents=True, exist_ok=True)
    with open(stats_dir / f"{spacecraft_id}_norm_stats.json", "w") as f:
        json.dump(stats, f, indent=2)

    time_res = config["model"]["time_resolution_minutes"]
    input_steps = (config["model"]["input_hours"] * 60) // time_res
    horizon_steps = (6 * 60) // time_res
    stride_steps = horizon_steps

    norm_feat_cols = sorted([c for c in processed.columns if c.endswith("_norm")])
    norm_tgt_cols = ["x_gse_norm", "y_gse_norm", "z_gse_norm"]

    inputs, targets, window_times = [], [], []
    for _, seg in processed.groupby("segment_id"):
        if len(seg) < input_steps + horizon_steps:
            continue
        feats = seg[norm_feat_cols].values
        tgts = seg[norm_tgt_cols].values
        times = seg["time"].values
        for i in range(0, len(seg) - input_steps - horizon_steps, stride_steps):
            inputs.append(feats[i:i + input_steps])
            targets.append(tgts[i + input_steps:i + input_steps + horizon_steps])
            window_times.append(times[i + input_steps])

    if not inputs:
        log.warning(f"{spacecraft_id}: no valid windows (data too sparse)")
        return {}

    inputs = np.array(inputs, dtype=np.float32)
    targets = np.array(targets, dtype=np.float32)
    window_times = np.array(window_times)

    test_start = int(0.85 * len(inputs))
    test_inputs = inputs[test_start:]
    test_targets = targets[test_start:]
    test_times = window_times[test_start:]

    if len(test_inputs) == 0:
        log.warning(f"{spacecraft_id}: empty test set")
        return {}

    sw_path = Path("data/raw/solar_wind_2023-01-01_2025-12-31.parquet")
    solar_df = pd.read_parquet(sw_path)
    test_kp = assign_kp(test_times, solar_df)

    log.info(f"Test: {len(test_inputs)} windows | quiet={np.sum(test_kp <= 3)}, "
             f"active={np.sum((test_kp >= 4) & (test_kp <= 5))}, storm={np.sum(test_kp >= 6)}")

    # Solar wind windows for multimodal
    test_solar = None
    try:
        from src.data.preprocessing import SolarWindPreprocessor
        sw_proc = SolarWindPreprocessor()
        solar_processed = sw_proc.preprocess(solar_df)
        aligned = sw_proc.align_with_positions(solar_processed, processed)
        solar_norm_cols = sorted([c for c in aligned.columns
                                  if c.endswith("_norm") and c.split("_norm")[0] in
                                  ["bx_gse", "by_gse", "bz_gse", "flow_speed",
                                   "proton_density", "kp", "dst", "ae"]])
        if solar_norm_cols:
            solar_inputs = []
            for _, seg in aligned.groupby("segment_id"):
                if len(seg) < input_steps + horizon_steps:
                    continue
                sw_feats = seg[solar_norm_cols].values
                for i in range(0, len(seg) - input_steps - horizon_steps, stride_steps):
                    solar_inputs.append(sw_feats[i:i + input_steps])
            if len(solar_inputs) == len(inputs):
                test_solar = np.array(solar_inputs, dtype=np.float32)[test_start:]
            else:
                log.warning(f"Solar mismatch: {len(solar_inputs)} vs {len(inputs)}")
    except Exception as e:
        log.warning(f"Solar prep failed: {e}")

    conditions = {
        "all": np.ones(len(test_kp), dtype=bool),
        "quiet": test_kp <= 3,
        "active": (test_kp >= 4) & (test_kp <= 5),
        "storm": test_kp >= 6,
    }

    model_results = {}
    for model_type in ["lstm", "transformer", "multimodal"]:
        ckpt_path = CHECKPOINT_DIR / f"{model_type}_{spacecraft_id}_6h_best.pt"
        if not ckpt_path.exists():
            log.warning(f"No checkpoint: {ckpt_path}")
            continue

        log.info(f"  {model_type}")
        try:
            model = load_model(model_type, ckpt_path, input_dim=inputs.shape[-1])
        except Exception as e:
            log.error(f"  Load failed: {e}")
            model_results[model_type] = {c: None for c in conditions}
            continue

        cond_results = {}
        for cond_name, mask in conditions.items():
            n = mask.sum()
            if n == 0:
                cond_results[cond_name] = None
                continue

            batch_size = 64
            all_preds = []
            masked_inputs = test_inputs[mask]
            masked_solar = test_solar[mask] if test_solar is not None else None

            for b in range(0, n, batch_size):
                with torch.no_grad():
                    x = torch.from_numpy(masked_inputs[b:b+batch_size]).float().to(DEVICE)
                    if model_type == "multimodal":
                        if masked_solar is not None:
                            sw = torch.from_numpy(masked_solar[b:b+batch_size]).float().to(DEVICE)
                        else:
                            sw = torch.zeros(x.shape[0], input_steps, 8).to(DEVICE)
                        p = model(x, sw)
                    else:
                        p = model(x)
                    all_preds.append(p.cpu().numpy())

            preds = np.concatenate(all_preds, axis=0)
            tgts = test_targets[mask]

            preds_km = np.zeros_like(preds)
            tgts_km = np.zeros_like(tgts)
            for i, col in enumerate(["x_gse", "y_gse", "z_gse"]):
                std = stats[spacecraft_id]["std"][col]
                mean = stats[spacecraft_id]["mean"][col]
                preds_km[..., i] = preds[..., i] * std + mean
                tgts_km[..., i] = tgts[..., i] * std + mean

            distances = np.sqrt(np.sum((preds_km - tgts_km)**2, axis=-1))
            valid = np.isfinite(distances)
            if valid.any():
                mae = round(float(np.nanmean(distances[valid])), 1)
            else:
                mae = None
            cond_results[cond_name] = mae
            log.info(f"    {cond_name} (n={n}): {mae:.1f} km")

        model_results[model_type] = cond_results
    return model_results


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--spacecraft", default="all")
    args = parser.parse_args()

    config = load_config()
    spacecraft_list = (
        list(config["spacecraft"].keys()) if args.spacecraft == "all"
        else [args.spacecraft]
    )

    all_results = {}
    for sc in spacecraft_list:
        sc_results = run_spacecraft(sc, config)
        for mt, cr in sc_results.items():
            if mt not in all_results:
                all_results[mt] = {}
            all_results[mt][sc] = cr

    RESULTS_DIR.mkdir(exist_ok=True)
    out_path = RESULTS_DIR / "storm_conditioned_mae.json"
    with open(out_path, "w") as f:
        json.dump(all_results, f, indent=2)
    log.info(f"Saved to {out_path}")


if __name__ == "__main__":
    main()