Getting started with specsr-roman

Roman’s grism will deliver spectra at R ≈ 461, where the diagnostic complexes stay blended: Hα+[N II], [O III]+Hβ. specsr-roman super-resolves them with a three-stage network, trained with losses that weight each line by how detectable it is in the grism data.

This notebook takes you from a fresh install to the numbers in the README:

1. download a small held-out subset of the training data,

2. load the published checkpoints,

3. super-resolve one spectrum and read what comes back,

4. see why the redshift is a PDF rather than a number,

5. measure what photometry buys, and why it must be noisy,

6. measure line recovery split by recoverability, for the lines the grism actually detects.

Everything runs on a laptop CPU in about two minutes and downloads roughly 17 MB (3.8 MB of data, 13 MB of weights).

Setup

pip install "specsr-roman[hub]"

The hub extra pulls in huggingface_hub, which is what fetches the published checkpoints and the tutorial data below.

import numpy as np
import torch
import matplotlib.pyplot as plt
from huggingface_hub import hf_hub_download

import specsr_roman
from specsr_roman import LINE_LIST_REST_AA, RomanPipeline
from specsr_roman.data import RomanFixedGridDataset, apply_phot_noise, normalize
from specsr_roman.evaluation import line_amplitude_recovery, redshift_summary
from specsr_roman.grids import ROMAN_MEDIUM_BANDS, WAVE_HR
from specsr_roman.lines import LABEL_LINES_AA
from specsr_roman.models import topk_modes

print("specsr-roman", specsr_roman.__version__, "| torch", torch.__version__)

# Everything here fits comfortably on a CPU; a GPU only makes it faster.
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", device)

specsr-roman 0.6.1 | torch 2.11.0+cu128
device: cuda

# Plot styling: recessive axes, one accent per model stage, truth always black.
INK, MUTED = "#0b0b0b", "#52514e"
LR_GREY, SR1_BLUE, SR2_ORANGE = "#9a9a94", "#2a78d6", "#eb6834"

plt.rcParams.update({
    "figure.dpi": 100, "savefig.dpi": 100,
    "font.size": 9, "axes.titlesize": 10, "axes.labelsize": 9,
    "axes.edgecolor": "#c9c8c2", "axes.labelcolor": MUTED,
    "axes.spines.top": False, "axes.spines.right": False,
    "xtick.color": MUTED, "ytick.color": MUTED,
    "grid.color": "#e8e7e2", "grid.linewidth": 0.8,
    "legend.frameon": False, "figure.facecolor": "white",
})

1. The tutorial dataset

The full training set is 180 MB and 23,845 spectra: every galaxy detected in the images whose Hα or [O III] flux exceeds the Roman survey’s line-flux limit. This notebook uses a 512-row subset of it, published alongside it on the Hub.

Two things about how that subset was drawn matter more than its size:

• It comes from the reserved test partition of the three-way galaxy split, so every number you compute below is an honest out-of-sample number. (The same galaxy appears in several visits as independent noise realisations, so a row-wise split would have measured memorisation instead.)

• It is a uniform random draw within that split — no cherry-picking on brightness or line strength. It carries the population’s real mix of detected and undetected lines. The redshift numbers below use all of it; the line recovery in section 6 is judged only on detected lines.

Rebuild it yourself with python scripts/make_tutorial_dataset.py.

path = hf_hub_download("aryana-haghjoo/romansr-data",
                       "tutorial/ou2024_h10307_tutorial.npz",
                       repo_type="dataset")
data = np.load(path, allow_pickle=True)

for key in ("flux_low", "flux_low_err", "flux_high", "redshift", "phot",
            "ids", "snr", "ab_h158", "wavelength_low", "wavelength_high"):
    print(f"{key:16s} {str(data[key].shape):12s} {data[key].dtype}")

flux_low         (512, 864)   float32
flux_low_err     (512, 864)   float32
flux_high        (512, 2500)  float32
redshift         (512,)       float64
phot             (512, 14)    float64
ids              (512,)       int64
snr              (512,)       float64
ab_h158          (512,)       float64
wavelength_low   (864,)       float64
wavelength_high  (2500,)      float64

Key	Meaning
flux_low, flux_low_err	the extracted grism spectrum and its 1σ error, 864 px on the native ~10.76 Å sampling
flux_high	the noiseless ground-truth SED, 2500 px — the super-resolution target
redshift	the true redshift
phot	catalogue fluxes; column order is fixed by grids.PHOT_BANDS
ids	OU2024 object_id — the thing to split on
snr, ab_h158	median extraction S/N and H158 magnitude

⚠️ flux_low and flux_high are not calibrated to each other — they differ by roughly twenty orders of magnitude. Only the shape is meaningful, which is why every stage of the pipeline normalises each spectrum individually.

z_all = data["redshift"]
print(f"{len(z_all)} spectra of {len(np.unique(data['ids']))} galaxies")
print(f"redshift   {z_all.min():.2f} - {z_all.max():.2f}  "
      f"(median {np.median(z_all):.2f}, {(z_all > 1).mean():.0%} above z = 1)")
print(f"AB(H158)   {data['ab_h158'].min():.1f} - {data['ab_h158'].max():.1f}")
print(f"median extraction S/N  {np.median(data['snr']):.2f}")

512 spectra of 434 galaxies
redshift   0.55 - 2.24  (median 0.93, 45% above z = 1)
AB(H158)   19.1 - 23.2
median extraction S/N  0.46

A median extraction S/N of about 0.4 is not a typo. These are single 301 s grism exposures, and most of these galaxies are faint. That is the regime the model has to be trustworthy in.

2. Load the published chain

from_pretrained() downloads the three canonical checkpoints once into the standard Hugging Face cache and reuses them afterwards.

from specsr_roman.checkpoints import CANONICAL_CHAIN

pipe = RomanPipeline.from_pretrained(device=device)
for stage, name in CANONICAL_CHAIN.items():
    print(f"{stage:6s} {name}")

sr1    sr1_ou2024_v6_independent
zhead  zhead_ou2024_roman_med3_noisy_independent
sr2    sr2_ou2024_v5_romanonly_independent

The three stages each do one job they can be held to:

• SR1 — a conservative residual CNN. It sees [flux, err] scaled by a shared factor, so the channel ratio literally is the per-pixel S/N and the network can matched-filter rather than guess which bumps are noise. It emits a mean and a log-variance.

• ZHead — a softmax over a redshift grid, conditioned on the coarse spectrum and the Roman imaging that ships with the grism. It returns a whole P(z).

• SR2 — attention with one token per rest-frame feature, run once per redshift hypothesis and combined by mode mass. That is what makes it robust when the top redshift mode is wrong.

3. Super-resolve one spectrum

RomanFixedGridDataset is what the training loops consume. We use it here for something the raw npz does not carry: the per-line integrated S/N labels, which record whether each line was recoverable from the low-resolution data at all. Those labels drive everything in section 6.

It also applies the dataset’s quality cuts, so index the raw arrays through ds.keep_mask before pairing them with anything the dataset returns — mixing the two indexings is a silent, and very confusing, off-by-N.

ds = RomanFixedGridDataset(path, with_phot=True, phot_tier="medium")

keep = ds.keep_mask
flux_low = data["flux_low"][keep]
flux_low_err = data["flux_low_err"][keep]
phot_true = data["phot"][keep][:, list(ROMAN_MEDIUM_BANDS)]
z_true = data["redshift"][keep]
ab_h158 = data["ab_h158"][keep]
wave_low = data["wavelength_low"]

print(f"{len(ds)} of {len(z_all)} rows pass the quality cuts")

quality cuts: kept 512 / 512 rows
phot tier 'medium': 3 bands ['roman_flux_Y106', 'roman_flux_J129', 'roman_flux_H158']
line-S/N labels: 45.1% of rows have a >3sigma recoverable line (median best-line S/N 2.77)
512 of 512 rows pass the quality cuts

phot_tier="medium" selects Roman F106/F129/F158 — and only those. That is the imaging that actually ships with the HLWAS grism, and feeding the redshift head anything more is how you accidentally measure the catalogue instead of the instrument.

# Recoverability of a row = the S/N of the best line the grism data actually shows.
best_line_snr = ds.line_snr.max(dim=1).values.numpy()

# An example with two lines to work with: the rows in which both H-alpha and
# [O III] 5007 are detected (columns 3 and 2 of the labels), and among them the
# one with the median H-alpha S/N. The rule looks only at the input, never at
# how well the model does.
snr = ds.line_snr.numpy()
two = np.where((snr[:, 3] >= 2) & (snr[:, 2] >= 2))[0]
i = int(two[np.argsort(snr[two, 3])[len(two) // 2]])
print(f"row {i}:  z = {z_true[i]:.3f},  AB(H158) = {ab_h158[i]:.1f},  "
      f"H-alpha S/N = {snr[i, 3]:.1f},  [O III] S/N = {snr[i, 2]:.1f}  "
      f"({len(two)} rows have both)")

row 458:  z = 1.268,  AB(H158) = 20.8,  H-alpha S/N = 7.9,  [O III] S/N = 2.1  (26 rows have both)

out = pipe.predict(
    flux_low[i],                              # (864,) on the native grism grid
    flux_low_err[i],
    phot=phot_true[i],                        # F106, F129, F158 — in that order
)

print(f"flux_sr      {out.flux_sr.shape}   super-resolved spectrum")
print(f"flux_sr_err  {out.flux_sr_err.shape}   per-pixel uncertainty")
print(f"flux_sr1     {out.flux_sr1.shape}   the SR1 stage alone")
print(f"pz / z_grid  {out.pz.shape}    the whole redshift PDF")
print(f"presence     {out.presence.shape}     per-line presence probability")
print()
print(f"z = {out.z:.4f} +/- {out.z_err:.4f}     (true {z_true[i]:.4f})")

flux_sr      (2500,)   super-resolved spectrum
flux_sr_err  (2500,)   per-pixel uncertainty
flux_sr1     (2500,)   the SR1 stage alone
pz / z_grid  (310,)    the whole redshift PDF
presence     (98,)     per-line presence probability

z = 1.2701 +/- 0.0311     (true 1.2684)

A note on flux scales

predict() returns values on your input’s flux scale, so they plot straight against the spectrum you handed it. The noiseless target lives on the simulation’s own scale instead, so to compare all four curves we put everything back in the shared normalised space the model works in.

Inverting that de-normalisation takes both the mean and the standard deviation. Dropping the mean leaves the output offset from the input by the continuum level — which looks exactly like a broken model and is not one.

def lr_scale(f_low):
    # The per-row (mean, std) predict() used to put its output on your scale.
    ok = np.isfinite(f_low)
    _, mean, std = normalize(np.interp(WAVE_HR, wave_low[ok], f_low[ok]))
    return mean, std


def to_model_space(flux, f_low):
    # Undo predict()'s de-normalisation: back to the shared normalised space.
    mean, std = lr_scale(f_low)
    return (flux - mean) / std


lr_n = ds[i][0][0].numpy()                    # LR input, on the HR grid, normalised
truth_n = ds[i][1].numpy()                    # noiseless target, normalised
sr1_n = to_model_space(out.flux_sr1, flux_low[i])
sr2_n = to_model_space(out.flux_sr, flux_low[i])
# An error bar takes the scale, not the offset.
sigma_n = out.flux_sr_err / lr_scale(flux_low[i])[1]

wave_um = WAVE_HR * 1e-4
peak = wave_um[np.argmax(truth_n)]
zoom = (peak - 0.05, peak + 0.05)


def mark_lines(ax, lo, hi):
    # Stagger the labels: at grism redshifts these complexes sit close together.
    for n, (name, rest) in enumerate(LABEL_LINES_AA):
        lam = rest * (1 + z_true[i]) * 1e-4
        if lo < lam < hi:
            ax.axvline(lam, color=MUTED, lw=0.7, ls=":", alpha=0.7, zorder=0)
            ax.annotate(name, (lam, 0.96 - 0.13 * (n % 2)),
                        xycoords=("data", "axes fraction"), ha="center",
                        va="top", fontsize=8, color=MUTED)


fig, axes = plt.subplots(3, 1, figsize=(9, 8.4), constrained_layout=True)
full = (wave_um >= 1.0) & (wave_um <= 1.93)

# 1. What the instrument gives you.
axes[0].plot(wave_um[full], lr_n[full], color=LR_GREY, lw=1.2)
axes[0].set_title(f"What the grism gives you  —  row {i}, z = {z_true[i]:.3f}",
                  color=INK, loc="left")

# 2. What the chain produces, against the noiseless truth.
axes[1].plot(wave_um[full], truth_n[full], color=INK, lw=1.5, label="truth")
axes[1].plot(wave_um[full], sr1_n[full], color=SR1_BLUE, lw=1.3, ls="--",
             label="SR1")
axes[1].plot(wave_um[full], sr2_n[full], color=SR2_ORANGE, lw=1.5, label="SR2")
axes[1].set_title("What the chain produces", color=INK, loc="left")
axes[1].legend(loc="upper center", ncols=3)

# 3. The strongest complex, with SR2's own uncertainty.
m = (wave_um >= zoom[0]) & (wave_um <= zoom[1])
axes[2].fill_between(wave_um[m], (sr2_n - sigma_n)[m], (sr2_n + sigma_n)[m],
                     color=SR2_ORANGE, alpha=0.20, lw=0, zorder=1,
                     label="SR2 $\\pm 1\\sigma$")
axes[2].plot(wave_um[m], lr_n[m], color=LR_GREY, lw=1.2, zorder=2,
             label="grism input")
axes[2].plot(wave_um[m], truth_n[m], color=INK, lw=1.6, zorder=5, label="truth")
axes[2].plot(wave_um[m], sr1_n[m], color=SR1_BLUE, lw=1.4, ls="--", zorder=3,
             label="SR1")
axes[2].plot(wave_um[m], sr2_n[m], color=SR2_ORANGE, lw=1.8, zorder=4,
             label="SR2")
axes[2].set_title("Zoom on the strongest complex", color=INK, loc="left")
axes[2].set_xlim(*zoom)
axes[2].legend(loc="upper left", ncols=2)
axes[2].set_xlabel("observed wavelength [µm]")

for ax, (lo, hi) in zip(axes, [(1.0, 1.93), (1.0, 1.93), zoom]):
    ax.set_xlim(lo, hi)
    ax.set_ylabel("normalised flux")
    ax.grid(axis="y", alpha=0.7)
    ax.set_axisbelow(True)
    mark_lines(ax, lo, hi)
plt.show()

The three panels have different y-scales, on purpose. Each spectrum is normalised by its own mean and standard deviation, and the input’s standard deviation is set by its noise while the target’s is set by its continuum. So a line that is a modest bump in panel 1 is a tall spike in panel 2 even though nothing was added — putting them on one axis would imply a calibration between them that does not exist.

Panel 2 shows the whole band. This galaxy has two detected lines, and the output draws both complexes: Hα with [N II] near 1.59 µm, and Hβ with the [O III] doublet near 1.2 µm, where the grism input shows little more than noise.

Read panel 3. SR1 (blue) is deliberately cautious and under-shoots the amplitude; SR2 (orange) sharpens it toward the truth and still falls a little short. Its peak also sits slightly to the red, because the inferred redshift is off by 0.0025. The shaded band is SR2’s own per-pixel uncertainty, and it is wide here — the model is telling you it knows the line is there but not exactly how bright. That is the intended behaviour: the alternative is a confident wrong amplitude.

What neither stage does is add lines elsewhere in the band. That restraint is the whole design, and section 6 is where it gets measured.

4. The redshift is a PDF, not a number

Redshift from a grism is a line-identification problem. A single observed line is consistent with Hα, [O III], [O II] or Lyα, and a model that returns one number has to average between those alternatives — which is how you get a confidently wrong answer. Two earlier designs of this head did exactly that, and both sat at a ~40 % catastrophic floor.

The ZHead returns a distribution over a grid instead. out.z is its mode; out.pz is the whole thing. Keep the PDF.

def top_modes(result, n=3):
    # The package's own mode finder -- the same hypotheses SR2 is run on.
    # It suppresses +/-15 grid bins around each pick, so the modes are distinct
    # line-alias candidates rather than adjacent bins of one peak, and weights
    # each by the probability *mass* in its window rather than by a bin height.
    zs, ws = topk_modes(torch.tensor(result.pz)[None, :],
                        torch.tensor(result.z_grid), n)
    return list(zip(zs[0].tolist(), ws[0].tolist()))


def plot_pz(result, truth, title):
    fig, ax = plt.subplots(figsize=(9, 3.6), constrained_layout=True)
    ax.fill_between(result.z_grid, result.pz, color=SR1_BLUE, alpha=0.30, lw=0)
    ax.plot(result.z_grid, result.pz, color=SR1_BLUE, lw=1.6, label="P(z)")
    ax.axvline(truth, color=INK, lw=1.4, ls="--", label="true redshift")
    # Stagger the mode labels by rank with a leader line: adjacent aliases sit
    # close enough in z that centred labels would overlap.
    for rank, (zm, mass) in enumerate(top_modes(result), start=1):
        height = result.pz[np.argmin(np.abs(result.z_grid - zm))]
        ax.plot([zm], [height], marker="o", ms=5, color=SR1_BLUE, zorder=3)
        ax.annotate(f"mode {rank} · {mass:.0%}", xy=(zm, height),
                    xytext=(zm, 1.05 - 0.09 * rank),
                    textcoords=("data", "axes fraction"), ha="center",
                    fontsize=8, color=MUTED,
                    arrowprops=dict(arrowstyle="-", lw=0.7, color="#c9c8c2"))
    ax.set_xlim(0, 2.6)
    ax.set_ylim(0, result.pz.max() * 1.5)     # headroom for the mode labels
    ax.set_xlabel("redshift")
    ax.set_ylabel("probability")
    ax.set_title(title, color=INK, loc="left")
    ax.legend(loc="upper right")
    ax.grid(axis="y", alpha=0.7)
    ax.set_axisbelow(True)
    plt.show()


for rank, (zm, mass) in enumerate(top_modes(out), start=1):
    print(f"mode {rank}:  z = {zm:.3f}   mass = {mass:.3f}")
print(f"\ntrue z = {z_true[i]:.3f}")

plot_pz(out, z_true[i], f"row {i}: unambiguous — one mode holds nearly all the mass")

mode 1:  z = 1.270   mass = 1.000
mode 2:  z = 1.978   mass = 0.000
mode 3:  z = 1.189   mass = 0.000

true z = 1.268

That galaxy is the easy case: one mode carries almost all the probability mass, it sits on the truth, and z_err is correspondingly small. The two remaining modes are the leftovers after suppression — real but negligible.

Section 5 finds the interesting case.

5. Photometry, and why it must be noisy

Three broadband colours break most of the alias degeneracy, which is the entire reason the redshift head takes photometry at all. But the catalogue photometry in this dataset is noiseless truth, and a metric measured on noiseless truth photometry is not a metric. apply_phot_noise adds the same 0.05 mag jitter the head was trained with.

We run the whole subset three ways: with realistic photometry, with the noiseless catalogue values, and with none.

gen = torch.Generator().manual_seed(0)         # seeded, so the run reproduces
phot_noisy = apply_phot_noise(torch.tensor(phot_true, dtype=torch.float32),
                              0.05, gen).numpy()


def run_chain(phot, batch=64):
    results = []
    for s in range(0, len(flux_low), batch):
        p = None if phot is None else phot[s:s + batch]
        results += pipe.predict(flux_low[s:s + batch],
                                flux_low_err[s:s + batch], phot=p)
    return results


outs = run_chain(phot_noisy)
outs_clean = run_chain(phot_true)
outs_nophot = run_chain(None)
print(f"{len(outs)} spectra through the chain, three ways")

512 spectra through the chain, three ways

configs = {
    "grism + Roman imaging, 0.05 mag noise": outs,
    "grism + noiseless catalogue photometry": outs_clean,
    "grism only (photometry mean-imputed)": outs_nophot,
}

print(f"{'configuration':40s} {'NMAD':>8s} {'catastrophic':>13s}")
for label, res in configs.items():
    m = redshift_summary(np.array([r.z for r in res]), z_true)
    print(f"{label:40s} {m['nmad']:8.4f} {m['catastrophic_frac']:12.1%}")

configuration                                NMAD  catastrophic
grism + Roman imaging, 0.05 mag noise      0.0012         0.6%
grism + noiseless catalogue photometry     0.0012         0.6%
grism only (photometry mean-imputed)       0.0013         3.1%

The noisy row is the deployable configuration. The full v0.6.0 reserved test set has NMAD 0.00133 and 0.67% catastrophic outliers on 2,384 spectra. This notebook uses a 512-row random draw, so its numbers will differ.

The noiseless row is a diagnostic using catalogue truth. It is not the survey performance estimate. phot=None supplies training-mean colours to a head trained with photometry; it is not a separately trained grism-only model.

fig, axes = plt.subplots(1, 2, figsize=(9, 4.4), constrained_layout=True,
                         sharex=True, sharey=True)
zline = np.linspace(0, 2.6, 200)

for ax, (label, res, colour) in zip(axes, [
        ("grism + Roman imaging (deployable)", outs, SR1_BLUE),
        ("photometry mean-imputed", outs_nophot, SR2_ORANGE)]):
    zp_i = np.array([r.z for r in res])
    frac = np.mean(np.abs(zp_i - z_true) / (1 + z_true) > 0.15)
    ax.fill_between(zline, zline - 0.15 * (1 + zline), zline + 0.15 * (1 + zline),
                    color="#e8e7e2", lw=0, zorder=0)
    ax.plot(zline, zline, color=INK, lw=1.0, zorder=1)
    ax.scatter(z_true, zp_i, s=9, color=colour, alpha=0.55, lw=0, zorder=2)
    ax.set_title(f"{label}\n{frac:.1%} catastrophic", color=INK, loc="left")
    ax.set_xlabel("true redshift")
    ax.set_xlim(0, 2.6)
    ax.set_ylim(0, 2.6)
    ax.grid(alpha=0.7)
    ax.set_axisbelow(True)

axes[0].set_ylabel("predicted redshift")
plt.show()

The shaded corridor is the ±0.15 catastrophic boundary. Compare both the core scatter and the outlier count: a tight core can coexist with ambiguous line identifications. The mean-imputed panel tests missing colours for this head; it does not establish the performance of a grism-only model.

Inspecting an ambiguous redshift

Inspect the widest PDF among the catastrophic outliers, or the widest PDF in the subset if it has no outliers. This is an illustrative example; the metrics above use every row.

zp = np.array([r.z for r in outs])
bad = np.where(np.abs(zp - z_true) / (1 + z_true) > 0.15)[0]
print(f"{len(bad)} catastrophic outliers out of {len(zp)}")
candidates = bad if len(bad) else np.arange(len(outs))
k = max(candidates, key=lambda r: outs[r].z_err)
modes = top_modes(outs[k])
print(f"row {k}: true z = {z_true[k]:.3f}, point estimate = {zp[k]:.3f} "
      f"(z_err = {outs[k].z_err:.3f})")
for rank, (zm, mass) in enumerate(modes, start=1):
    flag = "  <- near truth" if abs(zm - z_true[k]) / (1 + z_true[k]) < 0.05 else ""
    print(f"  mode {rank}: z = {zm:.3f} mass = {mass:.3f}{flag}")
plot_pz(outs[k], z_true[k], f"row {k}: inspect the full redshift PDF")

3 catastrophic outliers out of 512
row 238: true z = 0.678, point estimate = 1.131 (z_err = 0.196)
  mode 1: z = 1.131 mass = 0.782
  mode 2: z = 0.667 mass = 0.216  <- near truth
  mode 3: z = 0.765 mass = 0.001

Keep out.pz alongside the point estimate. Secondary modes expose competing line identifications that a single z cannot express. SR2 combines predictions across the top redshift hypotheses, weighted by mode mass. PDF width is not itself proof of calibrated uncertainty coverage.

6. Line recovery, split by recoverability

A single averaged “how much line flux did you recover” number mixes marginal detections with strong ones. So the score is binned by how well the best line is detected in the low-resolution data, using the per-line integrated S/N labels loaded in section 3:

bin	integrated line S/N
marginal	2 – 3
good	3 – 6
strong	> 6

The bins start at S/N 2. Below that a line is not a detection: the data carry no strong signal, so this metric does not judge the model there, in either direction. Inside a spectrum the same rule applies line by line: passing z and wave_um to line_amplitude_recovery counts only the lines that reach S/N 2 themselves. Unlike the redshift numbers above, which use every spectrum, this section therefore uses only the rows with a detected line.

truth_all = np.stack([ds[j][1].numpy() for j in range(len(ds))])
sr1_all = np.stack([to_model_space(r.flux_sr1, flux_low[j])
                    for j, r in enumerate(outs)])
sr2_all = np.stack([to_model_space(r.flux_sr, flux_low[j])
                    for j, r in enumerate(outs)])

rec = {stage: line_amplitude_recovery(pred, truth_all, ds.line_snr.numpy(),
                                      z=z_true, wave_um=WAVE_HR * 1e-4)
       for stage, pred in (("SR1", sr1_all), ("SR2", sr2_all))}

bins = list(rec["SR1"])
print(f"{'recoverability':16s} {'n':>5s} {'SR1':>8s} {'SR2':>8s}")
for b in bins:
    print(f"{b:16s} {rec['SR1'][b]['n']:5d} "
          f"{rec['SR1'][b]['median']:8.2f} {rec['SR2'][b]['median']:8.2f}")

recoverability       n      SR1      SR2
marginal           155     0.54     0.88
good               161     0.78     0.94
strong              70     0.79     0.88

fig, ax = plt.subplots(figsize=(7.4, 4.2), constrained_layout=True)
x = np.arange(len(bins))
w = 0.36

for off, stage, colour in ((-w / 2 - 0.01, "SR1", SR1_BLUE),
                           (w / 2 + 0.01, "SR2", SR2_ORANGE)):
    vals = [rec[stage][b]["median"] for b in bins]
    ax.bar(x + off, vals, w, color=colour, label=stage, zorder=2)
    for xi, v in zip(x + off, vals):
        ax.annotate(f"{v:.2f}", (xi, v), textcoords="offset points",
                    xytext=(0, 4 if v >= 0 else -12), ha="center",
                    fontsize=8, color=MUTED)

ax.axhline(1.0, color=INK, lw=1.0, ls="--", zorder=1)
ax.annotate("true amplitude", (len(bins) - 0.45, 1.0),
            textcoords="offset points", xytext=(0, 4), ha="right",
            fontsize=8, color=MUTED)
ax.axhline(0, color="#c9c8c2", lw=0.8, zorder=1)
ax.set_ylim(top=1.28)
ax.set_xticks(x)
ax.set_xticklabels([f"{b}\n(n = {rec['SR1'][b]['n']})" for b in bins])
ax.set_ylabel("median recovered flux fraction")
ax.set_title("Line amplitude recovery of the detected lines",
             color=INK, loc="left")
ax.legend(loc="upper left")
ax.grid(axis="y", alpha=0.7)
ax.set_axisbelow(True)
plt.show()

Compare the measured table with the full reserved-test Roman results: SR1 0.58 / 0.78 / 0.81 and SR2 0.90 / 0.94 / 0.89 for marginal / good / strong detections. Small bins in this 512-row subset fluctuate more. Flux recovery on simulated targets does not establish that individual output line fluxes are calibrated measurements.

Per-line presence probabilities

SR2 also emits a presence probability for each of its 98 rest-frame features. For the example galaxy from section 3:

rest_aa = dict(LINE_LIST_REST_AA)

print(f"{'feature':22s} {'p':>5s}   observed")
for j in np.argsort(out.presence)[::-1][:8]:
    name = out.line_names[j]
    lam_obs = rest_aa[name] * (1 + z_true[i])
    where = (f"{lam_obs * 1e-4:.3f} µm" if 10000 < lam_obs < 19300
             else "out of band")
    print(f"{name:22s} {out.presence[j]:5.2f}   {where}")

feature                    p   observed
[OIII]_5007             0.99   1.136 µm
Halpha_6563             0.96   1.489 µm
[NII]_6583              0.95   1.493 µm
[NII]_6548              0.94   1.485 µm
Hbeta_4861              0.92   1.103 µm
[OIII]_4959             0.90   1.125 µm
[OII]_3729              0.09   out of band
[OII]_3726              0.09   out of band

Features falling outside the grism band are gated off in SR2’s line branch, so a high presence probability there costs nothing and means nothing — read this list together with the observed wavelength.

7. From the command line

Everything above has a CLI equivalent, which is usually what you want for a catalogue rather than for a notebook.

!specsr-roman info

specsr-roman 0.6.1
torch 2.11.0+cu128 | cuda available: True
hub repo: aryana-haghjoo/roman-spectral-superresolution
canonical chain:
  sr1    sr1_ou2024_v6_independent
  zhead  zhead_ou2024_roman_med3_noisy_independent
  sr2    sr2_ou2024_v5_romanonly_independent
  extra [extract]: ok
  extra [train]: ok
  extra [hub]: ok

# run the chain over an npz holding flux_low[, flux_low_err, phot, wavelength_low]
specsr-roman predict spectra.npz --out predictions.npz --phot-tier medium

# reproduce the published metrics and figures from a frozen prediction cache
specsr-roman evaluate cache   --out pred_cache.npz
specsr-roman evaluate metrics --cache pred_cache.npz
specsr-roman evaluate figures --outdir figures/

# the audit that keeps the line-recovery claims honest
specsr-roman evaluate prior

Training each stage is one command and one config file:

specsr-roman train sr1   --config configs/sr1.yaml
specsr-roman train zhead --config configs/zhead.yaml     # needs SR1
specsr-roman train sr2   --config configs/sr2.yaml       # needs SR1 + ZHead

8. Where to go next

• Quickstart — the same API in reference form.

• Data guide — the full dataset, the splitting rule, and the photometric tiers.

• Training guide — including validation checkpoint selection (epoch 8 for the independent Roman SR2).

• Evaluation guide — the metrics above, plus the prior-dominance audit.

• ARCHITECTURE.md — the full design, the losses, and why each one is shaped the way it is.

Before you use this on anything real

• These results are on a simulation’s manifold. Targets are Diffsky model SEDs with that simulation’s line physics. A model can score well by learning the manifold rather than by measuring anything, and no reconstruction metric distinguishes the two — which is why the package ships a prior-dominance audit. The new Roman chain has deployed median response r ≈ 0.098, compared with 1 for proportional flux recovery. The central 68% P(z) interval covers 98.7% of reserved-test truths, so uncertainty calibration also remains necessary.

• Trained on simulations, not sky. Real Roman spectra will differ, and domain adaptation is open work.

• Line-identification aliases remain. Photometry helps disambiguate weak lines. Keep the PDF, but do not assume its intervals are calibrated or that the measured outlier rate is an irreducible physical floor.