Upload code/run_eval_296m.py

code/run_eval_296m.py

@@ -0,0 +1,415 @@
+"""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()