"""Frozen test-split predictions.
Every figure and every number quoted about this pipeline is computed from one
cache file rather than from a live model, for two reasons: the full chain over
the test split takes minutes on a GPU that is usually busy, and --- more
importantly --- a cached prediction set cannot drift. Regenerate it
deliberately when the chain changes; never regenerate it by accident while
tuning a plot.
"""
from __future__ import annotations
import os
from dataclasses import asdict, dataclass
import numpy as np
import torch
from ..checkpoints import CANONICAL_CHAIN, load_sr1, load_sr2, load_zhead_ckpt
from ..data import (
RomanFixedGridDataset,
apply_phot_noise,
evaluation_split,
get_or_make_group_split,
normalize,
)
from ..inference.pipeline import build_sr2_input
from ..lines import LINE_LIST_REST_AA, angstrom_to_micron
from ..models import constrain_delta
from .metrics import HALPHA_VAC_AA, OIII_VAC_AA, true_line_flux_cgs
__all__ = ["CacheConfig", "build_prediction_cache", "load_prediction_cache",
"add_true_line_fluxes"]
[docs]
@dataclass
class CacheConfig:
data: str = "data/dataset/ou2024_h10307_dataset.npz"
sr1_ckpt: str = CANONICAL_CHAIN["sr1"]
zhead_ckpt: str = CANONICAL_CHAIN["zhead"]
sr2_ckpt: str = CANONICAL_CHAIN["sr2"]
phot_tier: str | None = "medium"
#: Photometric noise applied at evaluation. Must match training, and must
#: not be zero: a metric measured on noiseless truth photometry is not a
#: metric. Seeded so the cache is reproducible.
eval_mag_err: float = 0.05
noise_seed: int = 0
batch_size: int = 64
delta_cap: float = 40.0
sigma_base_um: float = 0.005
z_topk: int = 3
#: False for a redshift head trained on photometry alone.
zhead_use_spectrum: bool = True
split_mode: str = "three_way"
partition: str = "test"
out: str = "outputs/pred_cache.npz"
def _true_line_fluxes(ds, rows) -> dict[str, np.ndarray]:
"""True H-alpha and [O III] 5007 fluxes [erg s^-1 cm^-2] for dataset ``rows``."""
if ds.ab_h158 is None:
raise KeyError("dataset has no `ab_h158` column, so its targets "
"cannot be put on an absolute flux scale")
args = (ds.wave_hi, ds.hi_raw[rows], ds.ab_h158[rows], ds.z[rows].numpy())
return {"ha_flux": true_line_flux_cgs(*args, HALPHA_VAC_AA),
"oiii_flux": true_line_flux_cgs(*args, OIII_VAC_AA)}
[docs]
def add_true_line_fluxes(path: str, data: str = CacheConfig.data) -> str:
"""Add ``ha_flux`` / ``oiii_flux`` to a cache written before they existed.
They depend only on the dataset and the split, so this runs no model and
leaves every prediction in the cache untouched.
"""
c = dict(np.load(path, allow_pickle=True))
ds = RomanFixedGridDataset(data, verbose=False)
_, test_idx, _ = get_or_make_group_split(os.path.abspath(data), ds.ids)
if (len(test_idx) != len(c["z_true"])
or not np.allclose(ds.z[test_idx].numpy(), c["z_true"])):
raise RuntimeError(f"{path} was not built from the test split of "
f"{data}; rebuild the cache instead")
c.update(_true_line_fluxes(ds, test_idx))
np.savez_compressed(path, **c)
print(f"added true line fluxes to {path}")
return path
[docs]
def build_prediction_cache(cfg: CacheConfig) -> str:
"""Run the full chain over the test split and cache every array a figure needs."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
zhead = load_zhead_ckpt(cfg.zhead_ckpt, device=device)
has_phot = getattr(zhead, "n_phot", 0) > 0
ds = RomanFixedGridDataset(cfg.data, with_phot=has_phot,
phot_tier=cfg.phot_tier if has_phot else None)
train_rows, test_idx, split_path = evaluation_split(
os.path.abspath(cfg.data), ds.ids, cfg.split_mode, cfg.partition)
from ..training.common import require_split_provenance
for ckpt in (cfg.sr1_ckpt, cfg.zhead_ckpt, cfg.sr2_ckpt):
require_split_provenance(ckpt, cfg.data, cfg.split_mode)
sr1 = load_sr1(cfg.sr1_ckpt, device=device)
wl_um = ds.wave_hi.astype(np.float32) * 1e-4
line_rest = angstrom_to_micron([w for _, w in LINE_LIST_REST_AA])
sr2 = load_sr2(cfg.sr2_ckpt, device=device, line_rest_um=line_rest,
wave_hi_um=wl_um)
# z normalisation from the TRAIN split, matching how the stages were trained.
ztr = ds.z[train_rows].numpy()
zmean, zstd = float(ztr.mean()), float(ztr.std())
zminn, zmaxn = (ztr.min() - zmean) / zstd, (ztr.max() - zmean) / zstd
run_cfg = {"z_topk": cfg.z_topk, "sigma_base_um": cfg.sigma_base_um,
"zhead_use_spectrum": cfg.zhead_use_spectrum}
cols: dict[str, list] = {k: [] for k in
("lr", "lr_err", "sr1", "sr2", "hr", "sigma",
"z_true", "z_pred", "line_snr",
"hr_mean", "hr_std", "lr_mean", "lr_std",
"ab_h158")}
pz_out = []
gen = torch.Generator(device=device).manual_seed(cfg.noise_seed)
with torch.no_grad():
for s in range(0, len(test_idx), cfg.batch_size):
rows = [int(i) for i in test_idx[s:s + cfg.batch_size]]
batch = [ds[i] for i in rows]
x_low = torch.stack([b[0] for b in batch]).to(device)
x_high = torch.stack([b[1] for b in batch])
z_true = torch.stack([b[3] for b in batch]).to(device)
snr = torch.stack([b[6] for b in batch])
phot = None
if has_phot:
phot = torch.stack([b[7] for b in batch]).to(device)
phot = apply_phot_noise(phot, cfg.eval_mag_err, gen,
zhead.phot_sigma)
sr1_out = sr1(x_low)
_, sr1_logvar = sr1_out
x_in, sr1_mean, z_modes, z_w, _ = build_sr2_input(
x_low, sr1, zhead, wl_um, line_rest, run_cfg, device,
phot=phot, z_mean=zmean, z_std=zstd,
z_min_n=zminn, z_max_n=zmaxn, sr1_out=sr1_out)
delta, logvar, _ = sr2(x_in, z_modes, z_w)
sr2_mean = sr1_mean + constrain_delta(delta, cfg.delta_cap)
z_in = torch.cat([x_low, sr1_mean, 0.5 * sr1_logvar], dim=1)
if not cfg.zhead_use_spectrum:
z_in = torch.zeros_like(z_in)
if hasattr(zhead, "z_centers"):
pz_out.append(torch.softmax(zhead(z_in, phot=phot), dim=-1).cpu().numpy())
cols["lr"].append(x_low[:, 0].cpu().numpy())
cols["lr_err"].append(x_low[:, 1].cpu().numpy())
cols["sr1"].append(sr1_mean[:, 0].cpu().numpy())
cols["sr2"].append(sr2_mean[:, 0].cpu().numpy())
cols["hr"].append(x_high.numpy())
cols["sigma"].append(
torch.exp(0.5 * logvar.clamp(-8, 6))[:, 0].cpu().numpy())
cols["z_true"].append(z_true.cpu().numpy())
cols["z_pred"].append(z_modes[:, 0].cpu().numpy())
cols["line_snr"].append(snr.numpy())
# What turns the normalised arrays back into flux: SR and HR share
# the target's statistics, the input has its own, and the H158
# magnitude is the only absolute scale a row carries.
cols["hr_mean"].append(np.array([float(b[4]) for b in batch]))
cols["hr_std"].append(np.array([float(b[5]) for b in batch]))
lr_stats = [normalize(ds.lo_raw[i])[1:] for i in rows]
cols["lr_mean"].append(np.array([m for m, _ in lr_stats]))
cols["lr_std"].append(np.array([s for _, s in lr_stats]))
cols["ab_h158"].append(np.asarray(ds.ab_h158[rows], dtype=np.float64))
print(f" cached {s + len(rows)}/{len(test_idx)}", end="\r")
print()
os.makedirs(os.path.dirname(os.path.abspath(cfg.out)), exist_ok=True)
import json
extra = ({"pz": np.concatenate(pz_out),
"z_centers": zhead.z_centers.cpu().numpy()} if pz_out else {})
np.savez_compressed(cfg.out, wl_um=wl_um,
object_id=ds.ids[test_idx], dataset_row=test_idx,
raw_row=np.flatnonzero(ds.keep_mask)[test_idx],
split_path=np.array(split_path),
config=np.array(json.dumps(asdict(cfg))),
**extra,
**{k: np.concatenate(v) for k, v in cols.items()},
**_true_line_fluxes(ds, test_idx))
print(f"wrote {cfg.out}")
return cfg.out
def load_prediction_cache(path: str, rebuild: bool = False,
cfg: CacheConfig | None = None):
if rebuild or not os.path.exists(path):
build_prediction_cache(cfg or CacheConfig(out=path))
elif "ha_flux" not in np.load(path, allow_pickle=True).files:
add_true_line_fluxes(path, (cfg or CacheConfig()).data)
return np.load(path, allow_pickle=True)