"""Line-perturbation sensitivity of frozen spectral reconstruction models.
Scale a recovered target line by a factor f, forward-model the difference
through a Gaussian grism LSF into the observed spectrum, and measure
r = log(L_perturbed / L_original) / log(f). The original errors and the
source's noisy photometry are held fixed.
The primary legacy_lr convention includes deployed input-dependent flux
restoration. A response of one is proportional output change; zero is no
change in that convention. target_fixed instead holds target statistics
constant and isolates normalized network output. Input normalization alone
can suppress its response even for an identity estimator, so it is not a
corrected physical response or a fraction attributable to a learned prior.
Summaries retain original-line and injected-change S/N breakdowns. Low
response to an injection below the noise can be appropriate. The Gaussian
spectral injection does not repeat two-dimensional morphological dispersion.
Per-response records support matched normalization controls and galaxy-grouped
bootstrap intervals. See the v0.6.1 audit report for the frozen test results.
"""
from __future__ import annotations
import os
from dataclasses import asdict, dataclass
import numpy as np
import torch
from ..checkpoints import load_sr1, load_sr2, load_zhead_ckpt
from ..data import (
RomanFixedGridDataset,
apply_phot_noise,
evaluation_split,
normalize,
)
from ..grids import GRISM_FWHM_AA
from .metrics import MIN_BEST_LINE_SNR, RECOVERABILITY_BINS
__all__ = ["PriorDominanceConfig", "run_prior_dominance", "plot_response"]
# np.trapz was renamed in NumPy 2.0; support both so the package does not
# force a NumPy major version on its users.
_trapz = getattr(np, "trapezoid", None) or np.trapz # noqa: NPY201
[docs]
@dataclass
class PriorDominanceConfig:
data: str = "data/dataset/ou2024_h10307_dataset.npz"
sr1_ckpt: str = "sr1_ou2024_v6_independent"
#: Sources enter with a detected line, at the threshold used everywhere
#: else (it was 5 until 2026-10-09, which left the marginal and most of the
#: good bin out of the audit).
snr_min: float = MIN_BEST_LINE_SNR
min_recovered_frac: float = 0.2 # line must be recovered at all
factors: tuple[float, ...] = (0.5, 2.0)
max_sources: int = 500
fwhm_aa: float = GRISM_FWHM_AA
#: Give both to audit the full pipeline as well as SR1. The redshift stage
#: re-reads the perturbed spectrum; the photometry is drawn once per source
#: and held fixed, so the response is to the spectrum alone.
zhead_ckpt: str | None = None
sr2_ckpt: str | None = None
#: Where to write the summary; the paper's macros and figure read it.
out_json: str | None = None
phot_tier: str = "medium"
eval_mag_err: float = 0.05
noise_seed: int = 0
delta_cap: float = 40.0
sigma_base_um: float = 0.005
z_topk: int = 3
zhead_use_spectrum: bool = True
split_mode: str = "three_way"
partition: str = "test"
# The target statistics are fixed across perturbations: changing the LR
# normalization must not itself amplify an unchanged standardized output.
output_scale: str = "legacy_lr" # production rescaling; target_fixed isolates normalized output
def _find_strongest_line(flux_hi, smooth_px: int = 101):
"""Contiguous segment around the strongest emission line, or ``None``."""
from scipy.ndimage import gaussian_filter1d
cont = gaussian_filter1d(flux_hi, smooth_px)
resid = flux_hi - cont
sigma = 1.4826 * np.median(np.abs(resid - np.median(resid)))
if sigma <= 0:
return None
mask = resid > 8 * sigma
if not mask.any():
return None
peak = int(np.argmax(resid))
if not mask[peak]:
return None
lo = hi = peak
while lo > 0 and resid[lo - 1] > 2 * sigma:
lo -= 1
while hi < len(resid) - 1 and resid[hi + 1] > 2 * sigma:
hi += 1
seg = np.zeros_like(mask)
seg[lo:hi + 1] = True
return seg, cont
def _line_flux(spec, seg, wave):
"""Continuum-subtracted line flux over a segment, local continuum removed."""
from scipy.ndimage import gaussian_filter1d
cont = gaussian_filter1d(spec, 101)
return _trapz((spec - cont)[seg], wave[seg])
def restore_prediction(pred, input_stats, target_stats, output_scale="target_fixed"):
"""Restore a normalized output in a declared, controlled unit system.
Training standardizes targets independently of LR inputs. Fixed original
target statistics make this a truth-conditioned sensitivity test, not an
inference-time absolute-flux calibration. The legacy convention allows an
LR scale change alone to change the measured output flux.
"""
if output_scale == "target_fixed":
mean, std = target_stats
elif output_scale == "legacy_lr":
mean, std = input_stats
else:
raise ValueError(f"unknown response output scale: {output_scale}")
return pred * std + mean
def recovery_is_usable(predicted, truth, scale, min_fraction):
"""Compare physical line integrals on the same scale, never mixed units."""
if not np.isfinite(scale) or scale <= 0 or not np.isfinite(truth) or truth <= 0:
return False
expected = scale * truth
return all(np.isfinite(v) and v > 0 and v >= min_fraction * expected
for v in predicted.values())
[docs]
def run_prior_dominance(cfg: PriorDominanceConfig) -> dict:
"""Run the audit on SR1, and on the full pipeline when it is configured.
Returns the response exponents per stage (``summary["sr1"]``,
``summary["sr2"]``), each per factor, overall and by recoverability bin. The
SR1 entries are repeated at the top level.
"""
if cfg.output_scale not in ("target_fixed", "legacy_lr"):
raise ValueError("output_scale must be target_fixed or legacy_lr")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sr1 = load_sr1(cfg.sr1_ckpt, device=device)
chain = bool(cfg.zhead_ckpt and cfg.sr2_ckpt)
zhead = load_zhead_ckpt(cfg.zhead_ckpt, device=device) if chain else None
has_phot = chain and 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_idx, 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):
if ckpt:
require_split_provenance(ckpt, cfg.data, cfg.split_mode)
with np.load(cfg.data) as raw:
native_wave = raw["wavelength_low"]
native_error = raw["flux_low_err"][ds.keep_mask]
wh = ds.wave_hi
sig_hr_px = (cfg.fwhm_aa / 2.355) / np.median(np.diff(wh))
# Candidates: bright enough that the line is genuinely in the data.
best_snr = ds.line_snr.max(dim=1).values.numpy()
cand = [i for i in test_idx if best_snr[i] >= cfg.snr_min][: cfg.max_sources]
print(f"{len(cand)} candidate test sources with a recoverable line")
if chain:
from ..inference.pipeline import build_sr2_input
from ..lines import LINE_LIST_REST_AA, angstrom_to_micron
from ..models import constrain_delta
wl_um = wh.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, as in the prediction cache
ztr = ds.z[train_idx].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}
gen = torch.Generator(device=device).manual_seed(cfg.noise_seed)
def run_models(lr_hr_grid, err_hr_grid, target_stats, phot=None):
"""Predictions in the fixed target or explicitly requested legacy units."""
xn, m, s = normalize(lr_hr_grid)
err_n = err_hr_grid / max(s, 1e-25)
x = torch.tensor(np.stack([xn, err_n])[None].astype(np.float32),
device=device)
with torch.no_grad():
if not chain:
pred, _ = sr1(x)
return {"sr1": restore_prediction(pred[0, 0].cpu().numpy(), (m, s),
target_stats, cfg.output_scale)}
x_in, sr1_mean, z_modes, z_w, _ = build_sr2_input(
x, 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)
delta, _, _ = sr2(x_in, z_modes, z_w)
sr2_mean = sr1_mean + constrain_delta(delta, cfg.delta_cap)
return {"sr1": restore_prediction(sr1_mean[0, 0].cpu().numpy(), (m, s),
target_stats, cfg.output_scale),
"sr2": restore_prediction(sr2_mean[0, 0].cpu().numpy(), (m, s),
target_stats, cfg.output_scale)}
from scipy.ndimage import gaussian_filter1d
def calibration_scale(hr, lr):
"""Scalar taking SED units into extraction units.
The Diffsky SEDs carry an internal flux scale (~1e-20 here) that has
nothing to do with the extraction's units, and the dataset hides the
mismatch by normalising each spectrum at load time. An injected delta
computed in SED units is therefore numerically invisible once added to
the LR spectrum -- which silently pins the response exponent at exactly
zero and makes every model look like it recites the prior.
The bridge is the least-squares scale between the LSF-smoothed truth
and the observation: the empirical flux-calibration ratio, absorbing
both the unit difference and the known ~1.7x aperture-loss offset.
"""
hr_s = gaussian_filter1d(hr, sig_hr_px)
ok = np.isfinite(hr_s) & np.isfinite(lr)
denom = float(np.dot(hr_s[ok], hr_s[ok]))
if denom <= 0:
return None
return float(np.dot(hr_s[ok], lr[ok])) / denom
stages = ("sr1", "sr2") if chain else ("sr1",)
# per stage: (factor, best-line S/N, r, source row) for every usable source
records: dict[str, list[tuple[float, float, float, float, float]]] = {
k: [] for k in stages}
n_used = 0
for candidate_number, i in enumerate(cand):
if candidate_number % 100 == 0:
print(f"audited {candidate_number}/{len(cand)} candidates", flush=True)
hr = np.nan_to_num(ds.hi_raw[i]).astype(np.float64)
lr = ds.lo_raw[i].astype(np.float64)
err = ds.err_raw[i].astype(np.float64)
found = _find_strongest_line(hr)
if found is None:
continue
seg, _ = found
phot = None
if has_phot:
phot = apply_phot_noise(torch.tensor(ds.phot[i])[None].to(device),
cfg.eval_mag_err, gen, zhead.phot_sigma)
_, target_mean, target_std = normalize(hr)
target_stats = (target_mean, target_std)
pred0 = run_models(lr, err, target_stats, phot)
L_pred0 = {k: _line_flux(v, seg, wh) for k, v in pred0.items()}
L_true0 = _line_flux(hr, seg, wh)
scale = calibration_scale(hr, lr)
if scale is None:
continue
# Fixed-target predictions and truth share units. Legacy predictions
# are in extraction units, so truth must be calibrated before gating.
gate_scale = 1.0 if cfg.output_scale == "target_fixed" else scale
if (not np.isfinite(scale) or scale <= 0
or not recovery_is_usable(L_pred0, L_true0, gate_scale,
cfg.min_recovered_frac)):
continue
cont = gaussian_filter1d(hr, 101)
for f in cfg.factors:
hr_p = hr.copy()
hr_p[seg] = cont[seg] + f * (hr[seg] - cont[seg])
# Forward-model the truth change into the observation, so the
# perturbed pair stays physically consistent -- LSF-smoothed to
# grism resolution and converted into the extraction's units.
delta = scale * gaussian_filter1d(hr_p - hr, sig_hr_px)
pred_p = run_models(lr + delta, err, target_stats, phot)
# Evaluate the injected change on native independent noise samples.
lo, hi = wh[seg][[0, -1]]
margin = 4 * cfg.fwhm_aa / 2.355
win = (native_wave >= lo - margin) & (native_wave <= hi + margin)
en = native_error[i, win]
valid_noise = len(en) > 0 and np.all(np.isfinite(en) & (en > 0))
injected_snr = (abs(np.interp(native_wave[win], wh, delta).sum())
/ np.sqrt(np.sum(en ** 2)) if valid_noise else np.nan)
for k in stages:
L_pred_p = _line_flux(pred_p[k], seg, wh)
if L_pred_p <= 0:
continue
r = np.log(L_pred_p / L_pred0[k]) / np.log(f)
if np.isfinite(r):
records[k].append((float(f), float(best_snr[i]), float(r),
float(i), float(injected_snr)))
n_used += 1
print(f"usable sources (line recovered above "
f"{cfg.min_recovered_frac:.0%} of truth): {n_used}")
summary: dict = {"n_sources": n_used, "n_candidates": len(cand),
"config": asdict(cfg), "split_path": split_path,
"method_version": 2,
"interpretation": "truth-conditioned fixed-scale sensitivity"
if cfg.output_scale == "target_fixed" else "LR-restored sensitivity"}
for k in stages:
rec = np.array(records[k]).reshape(-1, 5)
out: dict = {"per_factor": {}, "by_snr": []}
print(f"\n{k.upper()}")
for f in cfg.factors:
arr = rec[rec[:, 0] == f, 2]
if not arr.size:
continue
out["per_factor"][f] = {
"n": int(arr.size), "median": float(np.median(arr)),
"p25": float(np.percentile(arr, 25)),
"p75": float(np.percentile(arr, 75)),
}
print(f" f = {f}: N={arr.size} response exponent r: "
f"median {np.median(arr):.3f} "
f"p25/p75 = {np.percentile(arr, 25):.3f}/"
f"{np.percentile(arr, 75):.3f}")
out["overall_median_r"] = (float(np.median(rec[:, 2])) if len(rec)
else float("nan"))
print(f" OVERALL response exponent: {out['overall_median_r']:.3f} "
"(1 = proportional response, 0 = no response in the selected output scale)")
# The aggregate mixes lines the data constrain well with lines they
# barely constrain. Split by the recoverability bin of the best line.
for name, (lo, hi) in RECOVERABILITY_BINS.items():
m = (rec[:, 1] >= lo) & (rec[:, 1] < hi)
if m.sum() < 10:
continue
row = {"bin": name, "n": int(len(np.unique(rec[m, 3]))),
"median": float(np.median(rec[m, 2])),
"p25": float(np.percentile(rec[m, 2], 25)),
"p75": float(np.percentile(rec[m, 2], 75))}
out["by_snr"].append(row)
print(f" {name:9s} sources={row['n']:5d} median r {row['median']:.3f} "
f"p25/p75 = {row['p25']:.3f}/{row['p75']:.3f}")
out["by_injected_snr"] = []
for lo, hi in ((0, 1), (1, 3), (3, 6), (6, np.inf)):
mask = (rec[:, 4] >= lo) & (rec[:, 4] < hi)
if mask.any():
out["by_injected_snr"].append({
"lo": lo, "hi": hi if np.isfinite(hi) else None,
"n_responses": int(mask.sum()),
"n_sources": int(len(np.unique(rec[mask, 3]))),
"median": float(np.median(rec[mask, 2]))})
summary[k] = out
# the SR1 numbers stay at the top level, where earlier callers read them
summary.update(summary["sr1"])
if cfg.out_json:
import json
from pathlib import Path
os.makedirs(os.path.dirname(os.path.abspath(cfg.out_json)), exist_ok=True)
records_path = str(Path(cfg.out_json).with_suffix(".records.npz"))
np.savez_compressed(records_path,
columns=np.array(["factor", "best_line_snr", "response",
"dataset_row", "injected_snr"]),
**{k: np.asarray(v).reshape(-1, 5) for k, v in records.items()})
summary["records_path"] = records_path
with open(cfg.out_json, "w") as fh:
json.dump(summary, fh, indent=2)
print(f"wrote {cfg.out_json}")
return summary
[docs]
def plot_response(summary: dict, outdir: str = "outputs/figures",
name: str = "response.png") -> str:
"""Response exponent per recoverability bin, SR1 against the full pipeline.
Points are medians over the sources of a bin and bars span the quartiles.
The two reference lines are the two things a model can be doing: reading
the line strength from its input (1) or drawing the same line whatever the
input says (0).
"""
import matplotlib.pyplot as plt
from .figures import COLOR_LR, COLOR_SR, FigureStyle, _save
stages = [("sr1", "SR1", COLOR_LR, "o", -0.09),
("sr2", "SR2 (full pipeline)", COLOR_SR, "s", +0.09)]
with FigureStyle():
fig, ax = plt.subplots(figsize=(6.4, 4.6))
labels = None
for key, lab, col, mk, off in stages:
rows = summary.get(key, {}).get("by_snr", [])
if not rows:
continue
xs = np.arange(len(rows))
med = np.array([r["median"] for r in rows])
lo = med - np.array([r["p25"] for r in rows])
hi = np.array([r["p75"] for r in rows]) - med
ax.errorbar(xs + off, med, yerr=[lo, hi], color=col, marker=mk,
ms=8, lw=1.8, capsize=4, label=lab)
labels = [f"{r['bin']}\n$n$ = {r['n']:,}" for r in rows]
ax.axhline(1.0, color="0.3", lw=1.0, ls="--")
ax.axhline(0.0, color="0.3", lw=1.0, ls=":")
ax.text(0.99, 1.03, "proportional response", fontsize=9.5, color="0.3",
ha="right", transform=ax.get_yaxis_transform())
ax.text(0.99, -0.10, "no output response", fontsize=9.5, color="0.3",
ha="right", transform=ax.get_yaxis_transform())
if labels:
ax.set_xticks(np.arange(len(labels)), labels)
ax.set_xlim(-0.5, (len(labels) if labels else 3) - 0.5)
ax.set_ylim(-0.15, 1.25)
ax.set_xlabel("Line recoverability")
ax.set_ylabel(r"Response exponent $r$")
ax.legend(frameon=False, fontsize=10, loc="upper left",
bbox_to_anchor=(0.0, 0.93))
fig.tight_layout()
return _save(fig, name, outdir=outdir)