"""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()