code/run_eval_296m.py

"""Unified 75M evaluation: held-out DESI (Test 1), stress curve (Test 3), line-vs-continuum rec (Test 4).

Loads the 75M checkpoint once and runs many configurations against an external held-out cache.
"""
from __future__ import annotations

import argparse
import copy
import json
import math
from pathlib import Path
from typing import Any

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

from native_specz.data import SpectraListDataset
from native_specz.hybrid_redshift import HybridSpecZ, RawCollatorConfig, RawSpectraCollator, move_to_device
from native_specz.metrics import redshift_metrics


# ----- model + checkpoint -----

def build_model(ckpt: dict[str, Any], device: torch.device) -> HybridSpecZ:
    a = ckpt.get("args", {}) if isinstance(ckpt, dict) else {}
    model = HybridSpecZ(
        d_model=int(a.get("d_model", 256)),
        conv_width=int(a.get("conv_width", 128)),
        layers=int(a.get("layers", 5)),
        heads=int(a.get("heads", 8)),
        dropout=float(a.get("dropout", 0.1)),
        z_bins=int(a.get("z_bins", 64)),
        stem_stride=int(a.get("stem_stride", 8)),
        rec_hidden_mult=int(a.get("rec_hidden_mult", 0)),
        rec_refine_width=int(a.get("rec_refine_width", 16)),
        rec_refine_kernel=int(a.get("rec_refine_kernel", 5)),
        layerscale_init=float(a.get("layerscale_init", 0.0)),
        prediction_mode=str(a.get("prediction_mode", "regression")),
        bin_temperature=float(a.get("bin_temperature", 1.0)),
        residual_scale=float(a.get("residual_scale", 0.06)),
        candidate_topk=int(a.get("candidate_topk", 5)),
    ).to(device)
    state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
    missing, unexpected = model.load_state_dict(state, strict=False); print(f"NONSTRICT missing={len(missing)} unexpected={len(unexpected)}")
    model.eval()
    return model


# ----- perturbed sample helpers (for stress curve, Test 3) -----

def perturb_sample(sample: dict[str, Any], mode: str, strength: float, rng: np.random.Generator) -> dict[str, Any]:
    """Apply a real instrument-shift perturbation directly to a sample dict."""
    s = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in sample.items()}
    flux = s["flux"].astype(np.float32)
    ivar = s["ivar"].astype(np.float32)
    lam = s["lambda"].astype(np.float32)
    bad = s["bad_mask"].astype(np.bool_)
    n = len(flux)
    if mode == "wavelength_crop":
        # Keep a contiguous window covering fraction 1-strength of the spectrum.
        keep_frac = max(0.15, 1.0 - strength)
        width = max(64, int(n * keep_frac))
        start = int(rng.integers(0, max(1, n - width)))
        keep = np.zeros(n, dtype=np.bool_)
        keep[start : start + width] = True
        bad |= ~keep
    elif mode == "noise":
        # strength is a multiplier on typical sigma.
        good = np.isfinite(ivar) & (ivar > 0)
        sigma = np.zeros_like(flux)
        sigma[good] = 1.0 / np.sqrt(np.maximum(ivar[good], 1e-8))
        flux = flux + rng.normal(0.0, sigma * float(strength)).astype(np.float32)
    elif mode == "throughput":
        # strength scales the curvature amplitude.
        x = np.linspace(-1.0, 1.0, n, dtype=np.float32)
        amp = float(strength)
        coeff = rng.normal(0.0, [0.10 * amp, 0.05 * amp, 0.03 * amp]).astype(np.float32)
        curve = 1.0 + coeff[0] * x + coeff[1] * (x * x - 0.33) + coeff[2] * np.sin(np.pi * x)
        flux = flux * np.clip(curve, 0.3, 1.7).astype(np.float32)
    elif mode == "resolution":
        # Gaussian smoothing followed by replacement — simulates lower-resolution spectrograph.
        # strength = sigma in pixels.
        sigma_pix = max(1.0, float(strength))
        radius = max(3, int(math.ceil(3.0 * sigma_pix)))
        xs = np.arange(-radius, radius + 1, dtype=np.float32)
        k = np.exp(-0.5 * (xs / sigma_pix) ** 2)
        k = k / k.sum()
        good = np.isfinite(flux) & (~bad)
        f = np.where(good, flux, 0.0)
        w = good.astype(np.float32)
        f_sm = np.convolve(f, k, mode="same")
        w_sm = np.convolve(w, k, mode="same")
        flux = np.where(w_sm > 0.01, f_sm / np.maximum(w_sm, 1e-6), flux)
    else:
        raise ValueError(f"unknown mode {mode}")
    s["flux"] = flux.astype(np.float32)
    s["ivar"] = ivar.astype(np.float32)
    s["bad_mask"] = bad.astype(np.bool_)
    return s


class PerturbedDataset(torch.utils.data.Dataset):
    def __init__(self, samples: list[dict[str, Any]], mode: str, strength: float, seed: int):
        self.samples = samples
        self.mode = mode
        self.strength = float(strength)
        self.seed = int(seed)

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        s = self.samples[idx]
        if self.mode == "none" or self.strength <= 0:
            return s
        h = abs(hash((self.seed, self.mode, s["object_id"]))) % (2**32 - 1)
        rng = np.random.default_rng(h)
        return perturb_sample(s, self.mode, self.strength, rng)


# ----- core eval loop -----

@torch.no_grad()
def run_eval(
    model: HybridSpecZ,
    samples: list[dict[str, Any]],
    cfg: RawCollatorConfig,
    device: torch.device,
    *,
    perturb_mode: str = "none",
    perturb_strength: float = 0.0,
    batch_size: int = 16,
    num_workers: int = 2,
    collator_seed: int = 31415,
    max_samples: int | None = None,
) -> dict[str, np.ndarray]:
    if max_samples is not None and max_samples < len(samples):
        samples = samples[:max_samples]
    ds = PerturbedDataset(samples, perturb_mode, perturb_strength, seed=collator_seed)
    loader = DataLoader(
        ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
        collate_fn=RawSpectraCollator(cfg, train=False, seed=collator_seed),
    )
    z_true_l, y_true_l, y_pred_l, zwarn_l = [], [], [], []
    rec_l, rec_line_l, rec_cont_l = [], [], []
    line_count_l, cont_count_l = [], []
    oid_l: list[str] = []
    object_ids = [s["object_id"] for s in samples]
    idx_offset = 0
    for batch in tqdm(loader, desc=f"{perturb_mode}_s{perturb_strength:.2f}", leave=False):
        batch = move_to_device(batch, device)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            out = model(batch["x"], batch["valid"], batch["loglam"])
        y_pred = out.get("y_pred", out["y_mu"]).float()
        y_true = batch["y"].float()
        finite = torch.isfinite(y_true)
        z_true_l.append(batch["z"][finite].detach().cpu().numpy())
        y_true_l.append(y_true[finite].detach().cpu().numpy())
        y_pred_l.append(y_pred[finite].detach().cpu().numpy())
        zwarn_l.append(batch["zwarn"][finite].detach().cpu().numpy().astype(np.bool_))

        # Per-sample rec losses on masked pixels (all / line / continuum)
        rec = out.get("rec")
        bs = batch["y"].shape[0]
        if rec is not None and "target_flux" in batch and "loss_mask" in batch:
            per_pix = F.smooth_l1_loss(rec.float(), batch["target_flux"].float(), reduction="none", beta=0.5).detach().cpu().numpy()
            mask = batch["loss_mask"].detach().cpu().numpy().astype(np.float32)
            line_region = batch["line_region"].detach().cpu().numpy().astype(np.bool_)
            denom = mask.sum(axis=1).clip(min=1.0)
            rec_per = (per_pix * mask).sum(axis=1) / denom
            line_mask = mask * line_region.astype(np.float32)
            cont_mask = mask * (~line_region).astype(np.float32)
            line_denom = line_mask.sum(axis=1)
            cont_denom = cont_mask.sum(axis=1)
            rec_line_per = np.where(line_denom > 0, (per_pix * line_mask).sum(axis=1) / np.maximum(line_denom, 1.0), np.nan)
            rec_cont_per = np.where(cont_denom > 0, (per_pix * cont_mask).sum(axis=1) / np.maximum(cont_denom, 1.0), np.nan)
            rec_l.append(rec_per[finite.detach().cpu().numpy()])
            rec_line_l.append(rec_line_per[finite.detach().cpu().numpy()])
            rec_cont_l.append(rec_cont_per[finite.detach().cpu().numpy()])
            line_count_l.append(line_denom[finite.detach().cpu().numpy()])
            cont_count_l.append(cont_denom[finite.detach().cpu().numpy()])
        else:
            rec_l.append(np.full((int(finite.sum()),), np.nan, dtype=np.float32))
            rec_line_l.append(np.full((int(finite.sum()),), np.nan, dtype=np.float32))
            rec_cont_l.append(np.full((int(finite.sum()),), np.nan, dtype=np.float32))
            line_count_l.append(np.zeros((int(finite.sum()),), dtype=np.float32))
            cont_count_l.append(np.zeros((int(finite.sum()),), dtype=np.float32))

        finite_np = finite.detach().cpu().numpy()
        batch_oids = [object_ids[idx_offset + i] for i in range(bs)]
        oid_l.extend([o for o, ok in zip(batch_oids, finite_np) if ok])
        idx_offset += bs

    return {
        "z_true": np.concatenate(z_true_l).astype(np.float32),
        "y_true": np.concatenate(y_true_l).astype(np.float32),
        "y_pred": np.concatenate(y_pred_l).astype(np.float32),
        "zwarn": np.concatenate(zwarn_l).astype(np.bool_),
        "rec": np.concatenate(rec_l).astype(np.float32),
        "rec_line": np.concatenate(rec_line_l).astype(np.float32),
        "rec_cont": np.concatenate(rec_cont_l).astype(np.float32),
        "line_count": np.concatenate(line_count_l).astype(np.float32),
        "cont_count": np.concatenate(cont_count_l).astype(np.float32),
        "object_id": np.asarray(oid_l, dtype=object),
    }


def summarize(prefix: str, res: dict[str, np.ndarray]) -> dict[str, float]:
    y_true = res["y_true"]
    y_pred = res["y_pred"]
    metrics = {f"{prefix}/{k}": v for k, v in redshift_metrics(y_true, y_pred).items()}
    metrics[f"{prefix}/n"] = float(len(y_true))
    metrics[f"{prefix}/zwarn_fraction"] = float(np.mean(res["zwarn"])) if len(res["zwarn"]) else math.nan
    metrics[f"{prefix}/rec"] = float(np.nanmean(res["rec"])) if res["rec"].size else math.nan
    metrics[f"{prefix}/rec_line"] = float(np.nanmean(res["rec_line"])) if res["rec_line"].size else math.nan
    metrics[f"{prefix}/rec_cont"] = float(np.nanmean(res["rec_cont"])) if res["rec_cont"].size else math.nan
    metrics[f"{prefix}/rec_line_count_mean"] = float(np.nanmean(res["line_count"])) if res["line_count"].size else math.nan
    metrics[f"{prefix}/rec_cont_count_mean"] = float(np.nanmean(res["cont_count"])) if res["cont_count"].size else math.nan
    # Per-slice metrics
    z = np.expm1(y_true)
    slices = {"z_lt_0p4": z < 0.4, "z_0p4_1p0": (z >= 0.4) & (z < 1.0), "z_1p0_2p0": (z >= 1.0) & (z < 2.0), "z_gte_2p0": z >= 2.0}
    for name, mask in slices.items():
        if mask.sum() >= 5:
            sub = redshift_metrics(y_true[mask], y_pred[mask])
            for k, v in sub.items():
                metrics[f"{prefix}_slice/{name}/{k}"] = v
            metrics[f"{prefix}_slice/{name}/n"] = float(mask.sum())
    clean = ~res["zwarn"]
    if clean.any():
        sub = redshift_metrics(y_true[clean], y_pred[clean])
        for k, v in sub.items():
            metrics[f"{prefix}_clean/{k}"] = v
        metrics[f"{prefix}_clean/n"] = float(clean.sum())
        metrics[f"{prefix}_clean/rec"] = float(np.nanmean(res["rec"][clean])) if res["rec"][clean].size else math.nan
    return metrics


def ensemble_z_median(ys: list[np.ndarray]) -> np.ndarray:
    stack = np.stack(ys, axis=0)
    return np.nanmedian(stack, axis=0).astype(np.float32)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--cache", default="/workspace/native_specz_mae/cache/desi_heldout_2500.pt")
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--mask-ratios", default="0.30,0.50,0.65,0.75")
    parser.add_argument("--mask-mode", default="pixel")
    parser.add_argument("--mask-span-min", type=int, default=16)
    parser.add_argument("--mask-span-max", type=int, default=80)
    parser.add_argument("--tta-views", type=int, default=5)
    parser.add_argument("--max-samples", type=int, default=0)
    parser.add_argument("--skip-stress", action="store_true")
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print(f"LOAD_CHECKPOINT {args.checkpoint}")
    ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False)
    model = build_model(ckpt, device)
    ckpt_args = ckpt.get("args", {}) if isinstance(ckpt, dict) else {}
    target_length = int(ckpt_args.get("target_length", 8192))
    n_params = sum(p.numel() for p in model.parameters())
    print(f"MODEL_PARAMS {n_params}")

    cache_payload = torch.load(args.cache, map_location="cpu", weights_only=False)
    samples = cache_payload["samples"] if isinstance(cache_payload, dict) and "samples" in cache_payload else cache_payload
    if args.max_samples > 0:
        samples = samples[: args.max_samples]
    print(f"HELDOUT_SAMPLES {len(samples)}")

    mask_ratios = [float(x) for x in args.mask_ratios.split(",") if x]
    base_cfg = RawCollatorConfig(
        target_length=target_length,
        eval_mask_ratio=0.25,
        mask_mode=args.mask_mode,
        mask_span_min=args.mask_span_min,
        mask_span_max=args.mask_span_max,
        line_region_percentile=90.0,
    )

    all_metrics: dict[str, Any] = {
        "checkpoint": args.checkpoint,
        "cache": args.cache,
        "n_params": int(n_params),
        "heldout_n": int(len(samples)),
        "mask_ratios": mask_ratios,
        "mask_mode": args.mask_mode,
        "tta_views": int(args.tta_views),
    }

    # ===== Test 1a: held-out base eval at default mask 0.25 (no TTA, no aug) =====
    print("=== TEST 1a: held-out base eval (mask=0.25) ===")
    base_res = run_eval(model, samples, base_cfg, device, batch_size=args.batch_size, num_workers=args.num_workers, collator_seed=31415)
    all_metrics.update(summarize("heldout_base", base_res))
    np.savez_compressed(out_dir / "heldout_base.npz", **{k: v for k, v in base_res.items() if k != "object_id"}, object_id=base_res["object_id"].astype(str))

    # ===== Test 1b: TTA — multiple eval passes with different mask seeds + light aug =====
    print(f"=== TEST 1b: TTA ({args.tta_views} views) ===")
    tta_cfg = copy.deepcopy(base_cfg)
    tta_cfg.augment_ood = True
    tta_cfg.noise_prob = 0.4
    tta_cfg.throughput_prob = 0.4
    tta_views_y: list[np.ndarray] = [base_res["y_pred"]]  # include base prediction
    for v in range(args.tta_views):
        seed = 1000 + v * 17
        view_res = run_eval(model, samples, tta_cfg, device, batch_size=args.batch_size, num_workers=args.num_workers, collator_seed=seed)
        tta_views_y.append(view_res["y_pred"])
    # Align by object_id (they should match because data order is stable)
    y_tta_med = ensemble_z_median(tta_views_y)
    tta_res = dict(base_res)
    tta_res["y_pred"] = y_tta_med
    all_metrics.update(summarize("heldout_tta", tta_res))
    np.savez_compressed(out_dir / "heldout_tta.npz", y_pred_med=y_tta_med, y_pred_views=np.stack(tta_views_y, axis=0).astype(np.float32))

    # ===== Test 1c: multi-mask reconstruction sweep =====
    print("=== TEST 1c: multi-mask rec sweep ===")
    multi_mask: dict[str, Any] = {}
    for r in mask_ratios:
        cfg = copy.deepcopy(base_cfg)
        cfg.eval_mask_ratio = float(r)
        res = run_eval(model, samples, cfg, device, batch_size=args.batch_size, num_workers=args.num_workers, collator_seed=31415, max_samples=min(1000, len(samples)))
        key = f"heldout_mask{int(round(r*100)):02d}"
        sub = summarize(key, res)
        all_metrics.update(sub)
        multi_mask[key] = {"rec": sub[f"{key}/rec"], "rec_line": sub[f"{key}/rec_line"], "rec_cont": sub[f"{key}/rec_cont"], "n": sub[f"{key}/n"]}
    (out_dir / "multi_mask.json").write_text(json.dumps(multi_mask, indent=2), encoding="utf-8")

    # ===== Test 4: line-region vs continuum-region rec — already produced by every eval =====
    # We pull out the mask=0.50 case as the headline.
    print("=== TEST 4: line vs continuum rec at mask=0.50 ===")
    cfg = copy.deepcopy(base_cfg)
    cfg.eval_mask_ratio = 0.50
    cfg.mask_mode = "line_span"
    line_res = run_eval(model, samples, cfg, device, batch_size=args.batch_size, num_workers=args.num_workers, collator_seed=31415, max_samples=min(1500, len(samples)))
    line_metrics = summarize("heldout_linevscont50", line_res)
    all_metrics.update(line_metrics)
    np.savez_compressed(out_dir / "linevscont50.npz", **{k: v for k, v in line_res.items() if k != "object_id"})

    # ===== Test 3: stress curve =====
    if not args.skip_stress:
        print("=== TEST 3: stress curve ===")
        stress_results: dict[str, Any] = {}
        sweeps = {
            "wavelength_crop": [0.0, 0.20, 0.35, 0.50, 0.65],
            "noise": [0.0, 2.0, 5.0, 10.0],
            "throughput": [0.0, 1.0, 2.0, 4.0],
            "resolution": [0.0, 1.5, 3.0, 6.0],
        }
        stress_cfg = copy.deepcopy(base_cfg)
        stress_cfg.eval_mask_ratio = 0.25
        stress_samples = samples[: min(600, len(samples))]
        for mode, strengths in sweeps.items():
            mode_metrics: list[dict[str, Any]] = []
            for st in strengths:
                pmode = "none" if st == 0 else mode
                res = run_eval(model, stress_samples, stress_cfg, device, perturb_mode=pmode, perturb_strength=float(st), batch_size=args.batch_size, num_workers=args.num_workers, collator_seed=31415)
                sub = summarize(f"stress_{mode}_s{st:.2f}", res)
                mode_metrics.append({"strength": float(st), **sub})
                all_metrics.update(sub)
            stress_results[mode] = mode_metrics
            print(f"STRESS_DONE {mode}")
        (out_dir / "stress_curve.json").write_text(json.dumps(stress_results, indent=2), encoding="utf-8")

        # Plot stress curves
        fig, axes = plt.subplots(2, 2, figsize=(12, 8))
        for ax, (mode, results) in zip(axes.flat, stress_results.items()):
            xs = [r["strength"] for r in results]
            mae_z = [r[f"stress_{mode}_s{r['strength']:.2f}/mae_z"] for r in results]
            cat = [r[f"stress_{mode}_s{r['strength']:.2f}/cat_0p01"] for r in results]
            ax.plot(xs, mae_z, "o-", label="MAE(z)")
            ax2 = ax.twinx()
            ax2.plot(xs, cat, "s--", color="tab:red", label="Cat>0.01")
            ax.set_title(f"stress: {mode}")
            ax.set_xlabel("strength")
            ax.set_ylabel("MAE(z)")
            ax2.set_ylabel("Cat>0.01")
            ax.grid(alpha=0.2)
        fig.tight_layout()
        fig.savefig(out_dir / "stress_curve.png", dpi=150)
        plt.close(fig)

    # Save final summary
    (out_dir / "summary.json").write_text(json.dumps(all_metrics, indent=2, sort_keys=True), encoding="utf-8")
    print(f"WROTE {out_dir/'summary.json'}")
    headline_keys = [
        "heldout_base/mae_z", "heldout_base/nmad", "heldout_base/cat_0p01", "heldout_base/rec", "heldout_base/n",
        "heldout_tta/mae_z", "heldout_tta/nmad", "heldout_tta/cat_0p01",
        "heldout_clean/mae_z", "heldout_clean/cat_0p01", "heldout_clean/n",
        "heldout_linevscont50/rec_line", "heldout_linevscont50/rec_cont", "heldout_linevscont50/rec",
    ]
    print("HEADLINE")
    for k in headline_keys:
        if k in all_metrics:
            print(f"  {k}: {all_metrics[k]:.5f}")


if __name__ == "__main__":
    main()