API reference

Inference

Running the three stages as one pipeline.

Two entry points:

build_sr2_input()

Assembles SR2’s input channels and redshift hypotheses from frozen SR1 and a ZHead. Shared by training, evaluation and inference so all three see byte-identical inputs — a mismatch here is the kind of bug that shows up only as “the published metric does not reproduce”.

RomanPipeline

The user-facing object: give it a low-resolution grism spectrum (and photometry, if you have it) and get back a super-resolved spectrum with an uncertainty and a redshift PDF.

specsr_roman.inference.pipeline.build_sr2_input(x_low, sr1, zhead, wave_hi_um, line_rest_um, cfg, device, phot=None, z_mean=0.0, z_std=1.0, z_min_n=-3.0, z_max_n=3.0, sr1_out=None)[source]

(x_in, sr1_mean, z_modes, z_weights, line_mask).

cfg needs z_topk and sigma_base_um. Pass sr1_out as a precomputed (mean, log_var) to reuse a forward pass the caller has already made rather than running SR1 again.

For the P(z) head, M = z_topk distinct modes are extracted and the mask channel is their mass-weighted union. Giving SR2 several hypotheses rather than one point estimate is what makes it robust to alias errors: when the top mode is wrong, the correct identification is usually the second or third, and the line branch still draws there at reduced weight. For a regression head, M = 1 and there is nothing to be robust with.

Parameters:
class specsr_roman.inference.pipeline.RomanPipeline(sr1, zhead, sr2, device='cpu', wave_lr=None, wave_hr=None, z_topk=3, sigma_base_um=0.005, delta_cap=40.0)[source]

SR1 -> ZHead -> SR2, ready to run on a single spectrum or a batch.

>>> pipe = RomanPipeline.from_pretrained()
>>> out = pipe.predict(flux_low, flux_low_err, phot=phot)

phot must be the same band set the ZHead was trained on — three Roman Medium-tier fluxes (F106, F129, F158) for the published checkpoint, in that order, in any consistent linear flux unit.

Passing None drops the colour prior and costs most of the redshift accuracy — 26% catastrophic outliers on the held-out split against 5% with the imaging — because a single in-band line is alias-degenerate. It is not a grism-only model, though: this head was trained with photometry, so it receives its training-mean colours rather than nothing, which is mean imputation on an out-of-distribution input.

Parameters:
classmethod from_pretrained(sr1=None, zhead=None, sr2=None, device='cpu', repo_id=None, **kwargs)[source]

Load the published chain (or named alternatives) from the Hub.

Parameters:

  • sr1 (str | None)

  • zhead (str | None)

  • sr2 (str | None)

  • device (str)

  • repo_id (str | None)

Return type:

RomanPipeline

predict(flux_low, flux_low_err=None, phot=None, wave_low=None)

Run the chain. Accepts one spectrum or a batch of them.

Return type:

PipelineResult | list[PipelineResult]

class specsr_roman.inference.pipeline.PipelineResult(wavelength, flux_sr, flux_sr_err, flux_sr1, z, z_err, pz, z_grid, presence, line_names)[source]

What the pipeline returns, in the caller’s own flux units.

flux_sr and flux_sr_err are de-normalised back onto the input’s scale, so they can be plotted against the input directly. pz is the full redshift PDF over z_grid — keep it: for an alias-ambiguous source the second mode is real information that z alone discards.

Parameters:

Models

SR1 — the coarse super-resolution stage.

A fully convolutional 1D ResNet mapping the (upsampled) LR grism spectrum to a mean and a heteroscedastic log-variance on the HR grid. Architecture is deliberately identical to the JWST/JADES version in specsr so the two can be compared and warm-started across instruments; only the input channel count differs (Roman feeds [flux, err], see below).

class specsr_roman.models.sr1.SuperRes1D(*args, **kwargs)[source]

LR spectrum -> (mean, log_var) on the same grid.

Parameters:

  • in_channels (int) – 2 for the canonical Roman models: [flux, err], both divided by the same per-row flux scale so their ratio is the per-pixel S/N. The network can then matched-filter rather than guess which bumps are noise — the single largest quality jump in the Roman port. 1 is the JWST-compatible flux-only input.

  • hidden_dim (int) – Swept on the JWST data and carried over unchanged (120 / 16 / ~0.024).

  • num_res_blocks (int) – Swept on the JWST data and carried over unchanged (120 / 16 / ~0.024).

  • dropout (float) – Swept on the JWST data and carried over unchanged (120 / 16 / ~0.024).

  • activation_fn (nn.Module | None)

Notes

Input and output are the same length: the LR spectrum is interpolated onto the HR grid before it reaches the model. A fully convolutional stack with no resampling layers cannot change the axis length, and keeping the grids aligned means the residual pred - input is meaningful pixel by pixel.

The log-variance head is initialised to a constant -2.0 with zero weights, so the model starts by predicting a uniform sigma ~ 0.37 and has to earn any structure in its uncertainty.

classmethod from_state_dict(state, hidden_dim=120, num_res_blocks=16)[source]

Build a model whose input width matches a checkpoint, then load it.

The 1- vs 2-channel input is the one architectural thing that changed across Roman SR1 generations, and every downstream stage has to load whichever it is handed. Read it off the stem weight rather than making the caller remember.

Parameters:

  • state (dict)

  • hidden_dim (int)

  • num_res_blocks (int)

Return type:

SuperRes1D

ZHead — the redshift stage.

Three architectures live here because the first two failed in instructive ways and the failures explain the third:

ZHead1D

The direct JWST port: a Gaussian regression head over SR1’s (mean, log_sigma). On Roman it collapsed toward the prior mean (corr 0.44, predicted spread 0.27 against a true 0.60). Translation- equivariant convolutions plus global pooling carry no absolute positional signal, and unlike the JWST prism — whose varying resolution makes line width a wavelength cue — Roman’s near-constant R offers nothing to substitute. A normalised wavelength ramp channel (use_position) restores it.

ZHeadAttn

Feature-driven attention pooling. The earlier heads pooled with weights derived from SR1’s smooth sigma channel, which averaged a ~5-pixel emission line over 2500 pixels. Here saliency comes from the flux features themselves and each head also pools the wavelength ramp, so the attended ramp value is a line centroid: the redshift read out explicitly.

ZHeadClf (canonical)

P(z) classification over a fixed redshift grid. Redshift from a grism is a line-identification problem and is intrinsically multimodal — one observed line is consistent with Ha, [OIII], [OII] or Lya. A single Gaussian must average between aliases, which is precisely what produced the ~40% catastrophic-outlier floor of the regression heads. A softmax over a grid carries the whole P(z): the point estimate is the mode, and ambiguity survives as PDF width instead of becoming a wrong answer.

class specsr_roman.models.zhead.ZHead1D(*args, **kwargs)[source]

SR1 (mean, log_sigma) -> (mu_z, log_var_z). Superseded by ZHeadClf; kept so v1/v2 checkpoints stay loadable.

Parameters:
class specsr_roman.models.zhead.ZHeadAttn(*args, **kwargs)[source]

Feature-driven attention pooling with an explicit centroid readout.

Parameters:
class specsr_roman.models.zhead.ZHeadClf(*args, **kwargs)[source]

P(z) over a fixed redshift grid. The canonical redshift stage.

Parameters:

  • centers (torch.Tensor) – Bin centres of the redshift grid, from make_z_grid(). Stored as a buffer so a checkpoint is self-describing — SR2 reads the grid off the ZHead it is handed rather than being told it.

  • in_channels (int) – 4 for the canonical models: [LR flux, LR err, SR1 mean, SR1 log-sigma]. The raw LR channels matter more than they look: SR1 is the conservative stage and recovers only a few percent of line flux, so a head reading its near-line-free output alone was starved. The LR spectrum still holds the line at its true observed wavelength.

  • n_phot (int) –

    Number of broadband fluxes fed to the photometry branch, or 0. Colours over 0.35–2.1 um break the single-line alias degeneracy — exactly the information the grism band lacks. Fluxes go in raw; the branch applies log10 then standardises with train-split statistics carried in the phot_mu / phot_sig buffers.

    Feed it only bands that ship with the grism (see specsr_roman.grids.ROMAN_MEDIUM_BANDS). Given a complete, noiseless SED the head can read the redshift off the photometry alone, which measures the catalogue rather than the instrument.

  • hidden_dim (int)

  • num_blocks (int)

  • dropout (float)

  • n_heads (int)

  • refine_window (int)

predict_z(x, phot=None, window=None)

logits -> (zhat, sigma_z).

Parameters:

  • x (torch.Tensor)

  • phot (torch.Tensor | None)

  • window (int | None)

Return type:

tuple[torch.Tensor, torch.Tensor]

classmethod from_state_dict(state, **kwargs)[source]

Rebuild the exact head a checkpoint describes.

Grid, photometry width and input channel count are all recoverable from the saved tensors, so nothing about the head has to be remembered by the caller. feat_net.0 sees one extra channel (the wavelength ramp) that the constructor adds itself, hence the -1.

Parameters:

state (dict)

Return type:

ZHeadClf

specsr_roman.models.zhead.make_z_grid(z_lo, z_hi, n_bins, device='cpu')[source]

Bin centres of a uniform redshift grid.

Parameters:
Return type:

torch.Tensor

specsr_roman.models.zhead.soft_labels(z, centers, label_sigma)[source]

Gaussian-smoothed target over the grid (ordinal-soft cross-entropy).

A one-hot target would make adjacent bins as wrong as a distant alias. Smoothing tells the head that being one bin off is nearly right, which is what lets the mode-local refinement below reach sub-bin accuracy.

Parameters:

  • z (torch.Tensor)

  • centers (torch.Tensor)

  • label_sigma (float)

Return type:

torch.Tensor

specsr_roman.models.zhead.pz_stats(probs, centers, window)[source]

Mode-based point estimate plus full-grid PDF width.

zhat is the probability-weighted mean within +/- window bins of the mode — sub-bin accuracy without ever averaging across a distant alias. sigma_z is the standard deviation of the whole P(z), so it inflates precisely when the head is torn between two line identifications. The asymmetry is deliberate: the estimate should be decisive, the error bar should be honest about the ambiguity.

Parameters:

  • probs (torch.Tensor)

  • centers (torch.Tensor)

  • window (int)

Return type:

tuple[torch.Tensor, torch.Tensor]

specsr_roman.models.zhead.z_metrics(z_pred, z_true)[source]

The three numbers every redshift result is judged on.

dz_nmad is the robust scatter of the non-outliers; catastrophic_frac is the fraction beyond the conventional |dz|/(1+z) > 0.15. Report both — a model can trade one for the other, and NMAD alone hides an alias problem.

Return type:

dict[str, float]

specsr_roman.models.zhead.load_zhead(state, **kwargs)[source]

Build whichever ZHead variant a checkpoint holds.

SR2 and the evaluation scripts accept any generation, so the dispatch lives here rather than being repeated at each call site.

Parameters:

state (dict)

Return type:

torch.nn.Module

SR2 — the line-token attention refiner.

SR1 produces a smooth, conservative reconstruction; the ZHead produces a redshift PDF. SR2 spends its whole capacity on the delta between the two and the truth, through two branches:

Line branch. One token per rest-frame feature in specsr_roman.lines.LINE_LIST_REST_AA. Each token reads a local window of the input at that line’s predicted observed position, the tokens cross-attend (so [OIII]4959/5007 and Ha/[NII] can agree on a consistent picture), and each decodes to a gated Gaussian — amplitude x presence — scattered back onto the wavelength axis as a sparse delta.

CNN branch. A residual stack for the continuum and everything between lines.

The line branch runs once per redshift hypothesis and the deltas are combined weighted by P(z) mode mass. That is what makes the stage alias-robust: the correct line identification is almost always among the top few modes even when the point estimate is a catastrophic outlier, so the right placement still gets drawn, just at reduced weight.

class specsr_roman.models.sr2.SR2Attention(*args, **kwargs)[source]

(input channels, z hypotheses) -> (delta, log_var, presence).

Parameters:

  • in_channels (int) – 6 for the canonical models: LR flux, LR err, SR1 mean, SR1 sigma, line-position mask, broadcast zhat.

  • line_rest_um – Rest-frame line list and the observed grid, both in microns. Held as buffers so a checkpoint carries the line list it was trained with.

  • wave_hi_um – Rest-frame line list and the observed grid, both in microns. Held as buffers so a checkpoint carries the line list it was trained with.

  • window_half (int) – Half-width in pixels of the window each line token reads. 25 px on the HR grid is ~93 A — several resolution elements, enough to see the local continuum a line sits on.

  • line_dim (int)

  • num_attn_heads (int)

  • num_attn_layers (int)

  • cnn_dim (int)

  • num_cnn_blocks (int)

  • dropout (float)

forward(x, zhat, z_weight=None)[source]

zhat is (B,) for a single hypothesis or (B, M) for the top-M P(z) modes with matching z_weight.

Always returns (delta, log_var, presence) — presence in eval too, because every evaluation of this stage asks what it thought was there.

Parameters:

  • x (torch.Tensor)

  • zhat (torch.Tensor)

  • z_weight (torch.Tensor | None)

specsr_roman.models.sr2.topk_modes(probs, centers, k, suppress=15, refine=8)[source]

Top-k distinct modes of P(z) -> (z (B, k), weight (B, k)).

Iterative argmax with +/- suppress-bin suppression, so the k hypotheses are separate line-alias candidates rather than adjacent bins of one peak. Each is refined to a mode-local weighted mean and weighted by the probability mass inside its suppression window.

Parameters:

  • probs (torch.Tensor)

  • centers (torch.Tensor)

  • k (int)

  • suppress (int)

  • refine (int)

specsr_roman.models.sr2.constrain_delta(delta, cap)[source]

Soft-clip the SR2 delta.

The cap must exceed the tallest line the model has to reach: normalised SED lines peak above 30, and an early Roman run capped at 3 made them literally unreachable, saturating the gradient into “predict flat”. 40 is the working value.

Parameters:

  • delta (torch.Tensor)

  • cap (float)

Return type:

torch.Tensor

specsr_roman.models.sr2.build_line_mask(wave_hi_um, zhat, line_rest_um, sigma_base_um=0.005)[source]

(B,) redshifts -> (B, 1, L) Gaussian line-position mask.

Parameters:

  • wave_hi_um (torch.Tensor)

  • zhat (torch.Tensor)

  • sigma_base_um (float)

Return type:

torch.Tensor

specsr_roman.models.sr2.line_profiles(z, wave_um, line_rest_um, sigma_um=0.005)[source]

(B, K, L) per-line Gaussian windows at the given redshifts.

Parameters:

  • z (torch.Tensor)

  • wave_um (torch.Tensor)

  • sigma_um (float)

Return type:

torch.Tensor

Data

The training dataset.

One class, because all three stages consume the same rows: SR1 needs the spectra, the ZHead needs the spectra plus photometry, SR2 needs both plus the per-line recoverability labels. Building them once and returning a wide tuple keeps the three stages exactly aligned on which row is which.

class specsr_roman.data.datasets.RomanFixedGridDataset(*args, **kwargs)[source]

Roman grism LR/HR pairs on the shared fixed grids.

Loads a dataset npz, applies quality cuts, interpolates flux_low onto the HR grid, normalises per spectrum, and precomputes the per-line recoverability labels the losses need.

Returns per item:

(low_2ch, high, high_err, z, high_mean, high_std, line_snr[, phot])

low_2ch is [flux, err], both divided by the same per-row flux scale so that the channel ratio is the per-pixel S/N. That shared scaling is the point — normalising the error channel separately would destroy the very quantity the network reads.

high_err is zeros: the targets are noiseless simulated SEDs. It is still carried through so the loss signature matches the JWST version, where the targets are real grating spectra with real errors.

Parameters:

  • npz_path (str) – Dataset built by specsr_roman.extraction.

  • min_finite_low (float) – Minimum finite fraction for a row to survive.

  • min_finite_high (float) – Minimum finite fraction for a row to survive.

  • augment (bool) – Apply SpectrumAugmentor on the fly. Train split only — build a second instance for validation.

  • with_phot (bool) – Return the photometry column as an eighth tuple element.

  • phot_tier (str | None) – Band subset (see specsr_roman.grids.resolve_phot_tier()). Applied at load time so the dataset and the model it feeds cannot disagree about the band count.

  • verbose (bool)

Train/test splits.

The only rule that matters here: split by object, never by row. The same OU2024 galaxy is observed in many visits, each an independent noise realisation of the same underlying SED. A row-wise split puts realisation A in train and realisation B in test, and the resulting “held-out” metric is measuring memorisation.

specsr_roman.data.splits.hash_file(path)[source]

MD5 of a whole file — the identity of a dataset build.

Parameters:

path (str)

Return type:

str

specsr_roman.data.splits.get_or_make_group_split(dataset_path, ids, train_frac=0.8, split_dir=None, verbose=True)[source]

Deterministic split keyed by object id.

Membership is a pure hash of the id, not a shuffled permutation. Two consequences, both deliberate:

  • every stage of the pipeline derives the same split from the same ids, without having to pass a file around;

  • the split is stable under dataset growth — adding SCAs or visits never moves an existing galaxy across the boundary, so a model trained on the old build can still be evaluated on the new test set.

A record file is written for provenance, but membership never depends on reading it back.

Parameters:

  • dataset_path (str)

  • train_frac (float)

  • split_dir (str | None)

  • verbose (bool)

specsr_roman.data.splits.get_or_make_split(dataset_path, n_rows, train_frac=0.8, seed=42, split_dir=None, verbose=True)[source]

Row-wise split, for datasets with no object ids (Wang2022 era).

Prefer get_or_make_group_split() whenever ids exist.

Parameters:
specsr_roman.data.splits.filter_split_min_lines(train_idx, test_idx, z_all, wave_hi_aa, min_lines, verbose=True)[source]

Drop split rows with fewer than min_lines strong lines in band.

Applied after the shared split, never before: filtering first would change which galaxies fall on which side and break the guarantee that all three stages see the same partition. min_lines=2 selects the line-pair-identifiable population (roughly z > 1, about 53% of the OU2024 set) — a fair sample to quote redshift performance on when photometry is not available to break the alias.

Parameters:
specsr_roman.data.splits.default_split_dir(dataset_path)[source]

<dataset dir>/splits. Keeps split records beside the data they index.

Parameters:

dataset_path (str)

Return type:

str

Resampling and normalisation primitives.

Small functions, but each encodes a correctness lesson that cost a debugging session, so they live in one place with the reasoning attached.

specsr_roman.data.transforms.normalize(x, eps=1e-25)[source]

Per-spectrum standardisation -> (normalised, mean, std).

Every spectrum is normalised individually rather than by a global scale. That is not just conditioning: the grizli extraction is systematically ~1.7x brighter than the input SED (aperture losses), so an absolute flux scale would teach the model a calibration error. Normalising per row removes it and makes the task purely about shape.

std is floored because a genuinely constant row would otherwise divide by zero — though such rows should never reach here; see specsr_roman.data.datasets.RomanFixedGridDataset().

Parameters:
specsr_roman.data.transforms.fluxconserve_resample(wave, flux, new_wave)[source]

Rebin via the cumulative integral, conserving integrated flux.

Required for the Diffsky SEDs, whose wavelength grid is adaptive: sub-Angstrom bins at the emission lines and very coarse sampling elsewhere. Point-interpolating that onto a uniform grid drops or duplicates line flux depending on where the bins land — an emission line can simply vanish. Integrating and differencing preserves it exactly.

Parameters:
Return type:

ndarray

specsr_roman.data.transforms.smooth_to_grism(wave_obs, flux, resolution=461.0)[source]

Degrade a spectrum to Roman grism resolution.

R(lambda) = 461 * lambda[um], so FWHM = lambda/R = 1/461 um = 21.7 A — near-constant in wavelength, hence a constant-sigma Gaussian in lambda rather than a varying kernel. (This constancy is also why the ZHead needs an explicit wavelength ramp: unlike the JWST prism, line width here says nothing about where in the band you are.)

Parameters:
Return type:

ndarray

specsr_roman.data.transforms.interp_ascending(new_wave, wave, flux, left=nan, right=nan)[source]

np.interp that refuses to be fooled by a descending input grid.

Roman’s DLDP_A_1 is negative, so grizli’s optimal_extract returns wavelengths in descending order. np.interp does not check, and for unsorted xp it silently returns the edge value everywhere — producing a perfectly flat “spectrum” that carries zero information and trains a model straight into the prior mean. This bug cost one full dataset build.

Parameters:
Return type:

ndarray

Anti-prior augmentation.

The training targets are model SEDs (Galacticus for Wang2022, Diffsky for OU2024). A network can score well on them by learning the manifold — fixed line ratios, one dust law, one star-formation history family — rather than by reading line strengths out of the data. That is the classic inverse-crime failure, and it is invisible to every reconstruction metric.

The audit that detects it (specsr_roman.evaluation.prior_dominance) injects an off-manifold change to a line, forward-models it into the LR input, and measures a response exponent r: 1 means the model read the change from the data, 0 means it recited the prior regardless. The unaugmented Roman SR1 scored r = 0.14.

These augmentations attack the regularities directly:

  • per-line strength jitter, independent per line, so line ratios vary;

  • Calzetti dust jitter, because the simulations use one fixed Av;

  • an LR-only smooth calibration tilt, teaching invariance to flux-calibration error rather than to physics.

Every HR perturbation is forward-modelled onto the LR input through a Gaussian LSF at grism resolution (width jittered to stand in for morphological broadening), so the pair stays physically consistent — the augmented input is what the augmented truth would actually have produced.

With augmentation, r rose to 0.51 at fixed detectability. It also suppressed absolute line recovery, which is why the canonical SR1 does not use it: the honest fix is a redesign that jitters only recoverable information, or a broader simulation. See the limitations discussion in the README.

specsr_roman.data.augment.calzetti_k(lam_rest_um)[source]

Calzetti (2000) attenuation curve k(lambda), Rv = 4.05.

Parameters:

lam_rest_um (ndarray)

Return type:

ndarray

specsr_roman.data.augment.find_line_segments(flux_hi, smooth_px=101, thresh_sig=4.0, grow=2)[source]

Label contiguous emission-line segments in a noiseless HR spectrum.

Returns (labels, n_segments, continuum). Segment labels let the augmentor rescale one line without touching its neighbours, which is what makes the line ratios vary rather than the overall line strength.

Parameters:
class specsr_roman.data.augment.SpectrumAugmentor(wave_hi, fwhm_A=21.7, line_jitter=(0.4, 2.5), dav_range=0.4, tilt_amp=0.1, lsf_width_jitter=(1.0, 2.0))[source]

On-the-fly (lr, hr) -> (lr', hr') augmentation in raw flux space.

Applied before normalisation, and only to the train split — validation must always see the unaugmented distribution or the metric drifts with the augmentation settings.

Parameters:

Broadband photometry handling for the redshift stage.

Two operations, both of which have been got wrong at least once in this project’s history and both of which change the headline number by an order of magnitude.

specsr_roman.data.photometry.select_bands(phot, tier)[source]

Keep only the bands that ship with the grism.

tier is "medium" (Roman F106/F129/F158 — what the HLWAS grism actually comes with), an explicit "8,9,11", or None to keep everything.

Using "all" on OU2024 means feeding LSST ugrizy plus all eight Roman bands: an effectively complete SED, from which the redshift can be read without the spectrum contributing anything. It is a valid diagnostic and an invalid model.

Parameters:
Return type:

tuple[ndarray, tuple[int, …] | None]

specsr_roman.data.photometry.apply_phot_noise(phot, mag_err, generator=None)[source]

Multiplicative log-normal flux error of mag_err magnitudes.

Catalogue photometry in a simulation is noiseless truth. Training on it teaches the head to trust colours far beyond what a real measurement supports, and — worse — evaluating on it reports an accuracy nobody will reproduce. Apply this at train and validation both; pass a seeded generator for validation so checkpoint selection is not comparing epochs across different noise draws.

Parameters:

  • phot (torch.Tensor)

  • mag_err (float)

  • generator (torch.Generator | None)

Return type:

torch.Tensor

specsr_roman.data.photometry.band_names(indices)[source]

Index list -> OU2024 column names, for logging and model cards.

Return type:

list[str]

specsr_roman.data.photometry.standardization_stats(phot_train)[source]

log10 mean/std from the TRAIN split only -> ZHead buffers.

Computing these over the full set leaks test-set information into the input normalisation. It is a small leak next to feeding the whole SED, but it is free to avoid.

Parameters:

phot_train (ndarray)

Return type:

tuple[ndarray, ndarray]

Training

Loss functions.

The losses are where the Roman port differs most from its JWST ancestor, and the differences are not cosmetic — they follow from one fact about the data: the targets are noiseless simulated SEDs.

The JWST version got several scales for free from target noise. A robust MAD over the target’s derivative was a meaningful normaliser because the target had noise; the high-pass magnitude of a real grating spectrum was O(1) because it had noise. Feed those same terms a noiseless SED and the denominators collapse toward zero, and terms that were balanced against the NLL at O(1) arrive at 60x or 3500x its magnitude. Each such term is rescaled here by a quantity measured inside the line regions, which is O(1) in either regime.

The second theme is recoverability. About two thirds of Roman grism rows carry no line the data could possibly reveal. Under a plain reconstruction loss the optimal policy on those rows is to hedge — draw nothing, predict the prior mean — and that policy, averaged over the majority of the training set, is what SR1 converged to before the weighting below was introduced. Two terms express the fix: weight the reward for drawing a line by whether it is detectable, and separately penalise drawn amplitude where it is not.

specsr_roman.training.losses.make_line_mask_from_smoothed(x_high_raw, smooth_k=121, thresh_mad=7.5, dilate=11, min_width=7)[source]

Data-driven line mask: high-pass, threshold at thresh_mad MAD, clean up.

Derived from the target rather than from a redshift, so it is available even when the redshift is not, and it flags whatever structure is genuinely there.

Parameters:

  • x_high_raw (Tensor)

  • smooth_k (int)

  • thresh_mad (float)

  • dilate (int)

  • min_width (int)

Return type:

Tensor

specsr_roman.training.losses.line_flux_loss_weighted(mean, x_high, z_true, wave_um, line_rest_um, line_snr, sigma_um=0.005, floor=5.0, snr0=2.0, present_thresh=5.0)[source]

Integrated per-line flux L1 at true line positions, weighted by recoverability.

Plain reconstruction losses saturate: a 30-sigma amplitude miss and a 3-sigma one look similar once averaged over 2500 pixels, and the NLL can always be bought off by inflating the predicted variance at line pixels. Integrating the residual over each line window gives an unsaturated, per-line gradient that actually pushes a drawn line from 15% to 100% amplitude.

The weighting is what stops the hedging failure. Each present line is weighted snr/(snr + snr0), so lines the data cannot support contribute almost no gradient; absent lines keep weight 1, so drawing flux where there is none stays fully penalised. Asymmetric on purpose: we are permissive about failing to find the invisible, and strict about inventing.

Parameters:
Return type:

Tensor

specsr_roman.training.losses.line_hallucination_loss(mean, z_true, wave_um, line_rest_um, line_snr, sigma_um=0.005, floor=5.0, snr_h0=1.0, smooth_k=101)[source]

Penalise drawn line amplitude the data cannot support.

Recoverability weighting alone was not enough: given 300 epochs, SR1 still learned to recite the prior, drawing ~23% amplitude on lines that were present in the target but had integrated LR S/N below 1 — indistinguishable, to a user, from a real weak detection. This term measures the model’s own drawn flux (high-pass of the prediction, integrated over each line window) and penalises positive bumps by a non-recoverability weight.

The weight’s shape matters more than its size. A first attempt used the complement of the recovery weight, snr0/(snr + snr0), which is still 0.25 at S/N 6 — a quarter-strength drag on exactly the lines we want drawn — and it crushed recoverable recovery from 0.75 to 0.20. The squared knee used here, snr_h0^2 / (snr^2 + snr_h0^2), is 0.5 at snr_h0, 0.1 at 3x and 0.03 at 6x: recoverable lines are left alone while undetectable ones are driven back to the continuum.

Only emission (positive) bumps are penalised, so real absorption is safe. One caveat: high-passing a deep absorption line leaves positive wings, and a neighbouring line’s window can fall inside one (within roughly smooth_k pixels). The induced penalty is around two orders of magnitude below the emission case, so it does not drive training, but it is not identically zero.

Parameters:
Return type:

Tensor

specsr_roman.training.losses.loss_deblend_gated(mean, log_var, x_high, x_high_err, logvar_reg=3.4e-06, mask_smooth_k=121, mask_thresh_mad=7.5, mask_dilate=11, mask_min_width=7, lam_d1=0.11, lam_d2=0.0102, gate_min_frac=0.015, gate_temp=0.05, score_w_recon=0.2, score_w_line=2.0, row_w=None, eps=1e-12)[source]

SR1’s main objective: Gaussian NLL plus a gated, line-masked sharpness term.

The sharpness term compares first and second derivatives inside the line mask, which is what teaches deblending — matching a blended complex’s shape rather than only its integral. It is gated by how much of the row is masked, so rows with no lines do not get pushed toward spurious structure, and optionally by row_w (recoverability), so rows whose best line is buried in noise stop teaching the term to prefer flat outputs.

Returns (total, components); the components dict is what the training loop logs.

Parameters:
specsr_roman.training.losses.line_flux_loss(sr2_mean, x_high, z_true, wave_um, line_rest_um, sigma_um=0.005, floor=5.0)[source]

Unweighted integrated per-line flux L1 at true line positions.

SR2’s version of the teacher term. Lines absent from the target integrate to ~0, so drawing flux there is penalised by the same expression — which is how the stage learned to stop misallocating flux to Mg b and [NI]. floor turns strong lines into a relative error while keeping the absent-line penalty finite.

Parameters:
Return type:

Tensor

specsr_roman.training.losses.presence_labels(x_high, prof, hp_k=101, thresh=3.0)[source]

Ground-truth per-line presence from the noiseless HR target.

We know which lines each target has, so the presence head is supervised with BCE against these labels rather than left to discover them under a sparsity prior — which it never did: presence collapsed to zero in two successive SR2 generations, and with it every line the stage was meant to draw.

Parameters:
Return type:

Tensor

specsr_roman.training.losses.sr2_reconstruction_loss(*, sr2_mean, sr2_logvar, x_high, x_high_err, line_mask, presence, lam_hp_in, lam_hp_out, hp_k, lam_sparse, var_floor=1e-08)[source]

SR2 NLL plus in-line and out-of-line high-pass matching.

Splitting the high-pass term by the line mask lets the two regions carry different weights: sharpen hard where lines are, stay quiet elsewhere.

Parameters:

SR1 training loop.

specsr_roman.training.sr1.train(cfg)[source]

Train SR1 and return a summary dict (best monitor value, paths).

The checkpoint monitor is worth reading closely — it is not the validation loss. Total val loss is dominated by the NLL, which rises as the model commits to line amplitudes: a confident half-amplitude line costs more than a hedged flat one. An early run had every line-recovery metric improving to epoch 150 while the loss picked a mid-run “best”. The monitor here is the line-flux L1 (recovery) plus the hallucination penalty (so a checkpoint that buys recovery by reciting the prior is rejected), with residual RMS as a continuum tiebreak, smoothed by an EMA so a single lucky epoch cannot win.

Parameters:

cfg (SR1Config)

Return type:

dict

ZHead training loop.

specsr_roman.training.zhead.train(cfg)[source]

Train the redshift head against a frozen SR1.

Checkpoint selection is on NMAD, not the loss. On Roman the cross-entropy bottoms out early — driven by how well the PDF width is calibrated — while the point estimates keep improving for another hundred epochs. Selecting on the loss ships a worse model that is better at saying how unsure it is.

Parameters:

cfg (ZHeadConfig)

Return type:

dict

SR2 training loop.

specsr_roman.training.sr2.train(cfg)[source]

Train SR2 on top of a frozen SR1 and ZHead.

Checkpoint selection deserves the same warning as SR1’s, only more so. Plain validation NLL is continuum-dominated and reliably selects the SR2 that draws nothing. An earlier goal based on the line-region MSE ratio was no better: diluted across 98 line windows it is blind to sharpening, and it picked the most timid epoch available (strong-line amplitude 0.60, worse than the SR1 it was meant to improve).

The goal used here is an amplitude metric — -recov_amp + lam_hallu * hallu_amp — integrated predicted-over-true flux on recoverable strong lines against the same quantity on undetectable ones. It selects the sharpest epoch that is not yet hallucinating, which on the published run is epoch 4. An early best epoch here is the design working, not a truncated run: hallucination amplitude climbs from 0.26 to 0.62 by epoch 150 while recoverable amplitude barely moves.

Parameters:

cfg (SR2Config)

Return type:

dict

Evaluation

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.

class specsr_roman.evaluation.cache.CacheConfig(data: 'str' = 'data/dataset/ou2024_h10307_dataset.npz', sr1_ckpt: 'str' = 'sr1_ou2024_v6', zhead_ckpt: 'str' = 'zhead_ou2024_roman_med3_noisy', sr2_ckpt: 'str' = 'sr2_ou2024_v5_romanonly', phot_tier: 'str | None' = 'medium', 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, out: 'str' = 'outputs/pred_cache.npz')[source]
Parameters:
eval_mag_err: float = 0.05

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.

specsr_roman.evaluation.cache.build_prediction_cache(cfg)[source]

Run the full chain over the test split and cache every array a figure needs.

Parameters:

cfg (CacheConfig)

Return type:

str

Evaluation metrics.

The one thing to understand before reading any number this pipeline produces: line recovery must be reported split by recoverability. A single median amplitude ratio mixes lines the data clearly shows with lines it cannot possibly show, and a model can improve that average either by getting better or by hallucinating harder. The bins below separate the two, and the unrecoverable bin is the control: a well-behaved model scores near zero there. A model scoring 0.3 in that bin is inventing lines, however good its strong number looks.

Two amplitude metrics are provided because they answer different questions. line_amplitude_recovery() is the published one — per row, total predicted flux over total true flux across all line pixels, binned by the row’s best line. per_line_amplitude_recovery() scores each line separately in a Gaussian window, which is more diagnostic when you want to know which line a model gets wrong. They do not produce the same numbers and should not be compared to each other.

specsr_roman.evaluation.metrics.RECOVERABILITY_BINS: dict[str, tuple[float, float]] = {'good': (3.0, 6.0), 'marginal': (1.0, 3.0), 'strong': (6.0, inf), 'unrecoverable': (0.0, 1.0)}

Integrated line S/N in the LR input, and what each range means. Edges match the published evaluation; strong is open-ended above 6.

specsr_roman.evaluation.metrics.line_amplitude_recovery(pred, truth, line_snr, line_thresh=5.0, bins=None)[source]

Median recovered line-flux fraction per row, binned by recoverability.

This is the metric the published SR1 -> SR2 numbers are quoted from.

For each row, “line pixels” are those where the noiseless target exceeds line_thresh in normalised flux. The score is the summed predicted flux over the summed true flux across those pixels, which measures whether the line complex carries the right total amplitude without being fooled by a sub-pixel centroid error. Rows are binned by their best line’s integrated S/N, because a row’s recoverability is set by the line the data actually shows.

pred and truth are (N, L); line_snr is (N, K).

Parameters:
Return type:

dict[str, dict]

specsr_roman.evaluation.metrics.per_line_amplitude_recovery(pred, truth, line_snr, z, wave_um, line_rest_um=None, sigma_um=0.005, present_thresh=3.0, bins=None)[source]

Median integrated predicted/true flux per line, binned by that line’s S/N.

Diagnostic companion to line_amplitude_recovery(): it attributes a failure to a specific transition rather than to a row. Only lines actually present in the target are scored — absent lines are the hallucination test, which the unrecoverable bin already covers.

Parameters:
Return type:

dict[str, dict]

specsr_roman.evaluation.metrics.redshift_summary(z_pred, z_true)[source]

NMAD, median \(|\Delta z|/(1+z)\), catastrophic fraction, and N.

Report all of them. NMAD alone describes only the well-behaved core, and a model can shrink it while pushing more objects past the catastrophic threshold.

Return type:

dict[str, float]

specsr_roman.evaluation.metrics.line_snr_from_spectrum(wave_um, flux, lam_obs_um, half=0.045, core=0.012, sbgap=0.015, sbw=0.03)[source]

Local S/N of one line, measured off a spectrum using local sidebands.

Used to compare the S/N a line has in the LR input against the S/N it has after super-resolution — the “did this become measurable” question, which is distinct from “was its amplitude right”.

Parameters:
Return type:

float

Photometry ablation — how much of the redshift accuracy is the spectrum?

The redshift head takes two inputs: the grism spectrum (raw plus SR1’s reconstruction) and three Roman Medium-tier colours. A number from that head is only interpretable if you know which input produced it, and the way to find out is to take the photometry away.

Photometry enters the head standardised as (log10(flux) - phot_mu) / phot_sig, with the statistics baked into the checkpoint, so “remove the colours” is exactly “feed every band its training mean”, which standardises to zero. That is the same vector the head receives from RomanPipeline.predict() when it is called with phot=None, so the grism only row below is a measurement of the deployed model in that mode rather than of a hypothetical one.

Read the ``grism only`` row as an upper bound, not as an information floor. This head was trained with colours. Handing it a mean-imputed colour vector tells you what the deployed chain does when photometry is missing; it does not tell you how well a head trained without colours would do, because the two differ by everything the network learned to delegate to the photometry branch. A grism-only head is a separate experiment and has not been run. The physical floor is set by the alias degeneracy — with a single line in band, H-alpha, [O III] and [O II] are mutually consistent — and no architecture removes it.

Why there is no “zero the spectrum” row. The obvious complement — keep the photometry, blank the spectrum, see what the colours alone can do — does not work here, and reporting it would be worse than reporting nothing. Masking a band to its training mean is in distribution: it standardises to exactly 0, a value the network sees constantly. There is no equivalent for the spectral channels. Feeding SR1 a zero array produces a reconstruction and an uncertainty map unlike anything in training, and the head then reads nonsense from two of its four channels. Measured, that configuration scores worse than removing the photometry — which tells you the input was out of distribution, not what the photometry contributes. The grism only row answers the answerable half of the question; the other half needs a head trained without the spectrum.

The noise sweep is the quantitative version of the same question. It perturbs the three colours with the multiplicative log-normal jitter used in training, so the small levels are in distribution and the degradation is real rather than an artefact of an unfamiliar input. Each level is drawn from a generator reseeded to the same value, so the sweep varies only sigma.

Every configuration here uses the three Roman Medium-tier bands that ship with the HLWAS grism, and nothing else.

class specsr_roman.evaluation.ablation.AblationConfig(data: 'str' = 'data/dataset/ou2024_h10307_dataset.npz', sr1_ckpt: 'str' = 'sr1_ou2024_v6', zhead_ckpt: 'str' = 'zhead_ou2024_roman_med3_noisy', roman_bands: 'tuple[int, ...]' = (8, 9, 11), eval_mag_err: 'float' = 0.05, noise_seed: 'int' = 0, batch_size: 'int' = 64, out_dir: 'str' = 'outputs', noise_levels: 'tuple[float, ...]' = (0.0, 0.02, 0.05, 0.1, 0.2, 0.3, 0.5, 1.0))[source]
Parameters:
zhead_ckpt: str = 'zhead_ou2024_roman_med3_noisy'

The deployable Roman Medium-tier head. There is deliberately no wider-band head here: a model fed bands the survey does not deliver alongside the grism reads the redshift off an effectively complete SED and measures the simulation rather than the instrument.

specsr_roman.evaluation.ablation.run_ablation(cfg)[source]

Run the ablation and write phot_ablation.csv. Returns the rows.

Parameters:

cfg (AblationConfig)

Return type:

list[dict]

specsr_roman.evaluation.ablation.plot_ablation(rows, out_dir='outputs')[source]

Render phot_ablation.png: what the colours buy, and how fast.

Left, the two configurations that matter — the deployed chain and the same chain with its colours removed. Right, the noise sweep, which is the same question asked continuously: the operating point is marked, and the y-axis is the outlier rate because that, not the scatter, is what a survey pipeline pays for.

Parameters:
Return type:

str

The inverse-crime audit: is the model reading the data or reciting the prior?

The training targets are model SEDs. A network can score well on every reconstruction metric by learning the simulation’s manifold — its fixed line ratios, its single dust law — rather than by measuring anything. No reconstruction metric can tell the two apart, because on the manifold they give the same answer.

This test forces them apart. Take a source whose line the model does recover. Scale that line in the truth by a factor f — deliberately off the manifold, a line ratio the simulation never produces — forward-model the difference through a Gaussian LSF at grism resolution onto the observed LR spectrum, and re-run the model. Then measure the response exponent

r = log(L_pred_perturbed / L_pred_original) / log(f)

r = 1 means the model tracked the change: it read the line strength from the data. r = 0 means it produced the same line regardless: it recited the prior.

Reading the result requires care. Where the injected change is genuinely below the noise, a low r is the correct behaviour — falling back on the prior is what a well-calibrated model should do when the data says nothing. So bin by the detectability of the injected change, and judge r only where the information is physically present. Aggregate r is dominated by unrecoverable cases and understates a good model.

Measured on this project: unaugmented SR1 scored 0.14; with anti-prior augmentation, 0.51 at fixed detectability — a 3.7x improvement in data-faithfulness, at the cost of absolute line recovery, which is why the published SR1 is the unaugmented one and this remains open work.

class specsr_roman.evaluation.prior_dominance.PriorDominanceConfig(data: 'str' = 'data/dataset/ou2024_h10307_dataset.npz', sr1_ckpt: 'str' = 'sr1_ou2024_v6', snr_min: 'float' = 5.0, min_recovered_frac: 'float' = 0.2, factors: 'tuple[float, ...]' = (0.5, 2.0), max_sources: 'int' = 500, fwhm_aa: 'float' = 21.7)[source]
Parameters:
specsr_roman.evaluation.prior_dominance.run_prior_dominance(cfg)[source]

Run the audit. Returns per-factor and overall response exponents.

Parameters:

cfg (PriorDominanceConfig)

Return type:

dict

Extraction

Resumable batch driver: OU2024 products -> a training dataset.

Design constraints that shaped this file:

  • grizli’s numba disperser corrupts the heap when thousands of sources are dispersed in one process. Each (visit, SCA) therefore runs in its own subprocess, and a crash costs one SCA rather than the run.

  • Everything is resumable. Each SCA writes a cached npz — including a stub recording why it produced nothing — so re-running skips completed work and never silently retries a pointing that is outside the healpix.

  • The same galaxy appears in many visits. Rows carry ids, and the merge step ends by reminding you to split on them.

class specsr_roman.extraction.batch.ExtractionConfig(healpix=10307, data_dir='data/ou2024', out_dataset='data/dataset/ou2024_h10307_dataset.npz', max_scas=60, workers=4, ab_target=22.5, ab_scene=23.0, min_hp_frac=0.7, grism_exptime=301.0, cleanup=False, z_max=3.2)[source]

Where the inputs are, and what counts as a target.

Parameters:
specsr_roman.extraction.batch.run_worker(visit, sca, cfg)[source]

Extract every target on one (visit, SCA) and cache the rows.

Parameters:
Return type:

None

specsr_roman.extraction.batch.run_batch(cfg)[source]

Drive workers over candidate SCAs until max_scas produce rows.

Parameters:

cfg (ExtractionConfig)

Return type:

None

specsr_roman.extraction.batch.merge(cfg)[source]

Concatenate the per-SCA caches into one training dataset.

Parameters:

cfg (ExtractionConfig)

Return type:

str

Turning simulation products into grizli-ready frames.

OpenUniverse2024 ships direct images and truth catalogues but no grism images, so this package disperses the scene itself with grizli (the same approach as Guo et al. 2025, arXiv:2512.09993). That is not a workaround — it is what makes the training pairs self-consistent: the same configuration file disperses the scene and extracts it back, so any residual is the fault of noise, blending and the extraction, not of a mismatched instrument model.

What this module does is write the two FITS files grizli expects: a direct image in e-/s with a photometric calibration, and an empty grism shell sharing its WCS, tagged INSTRUME='WFI' so grizli picks up Roman.G150.conf.

specsr_roman.extraction.frames.prepare_frames(img_path, out_dir, grism_exptime=301.0, overwrite=False)[source]

OU2024 simple_model image -> (direct_fits, grism_fits).

The direct image is converted from counts to e-/s with the flat sky level removed, because grizli’s source model works in rate units and the zodiacal background is added back explicitly at the noise stage.

PHOTFLAM/PHOTPLAM are the Wang2022 F158 values, and they apply to OU2024 unchanged: those images are calibrated to the real Roman zeropoint (verified — an AB ~ 14 index galaxy realises 1.26e7 e- in 139.8 s, exactly 10^(-0.4 (14 - 26.4)) * 139.8).

A caveat that bites anything reading the pixels directly: setting PHOTFLAM makes grizli rescale the SCI array into f_lambda units internally. Downstream code must be unit-agnostic or read the header.

Parameters:
Return type:

tuple[str, str]

Truth catalogues, source detection, and the match between them.

Every extracted spectrum has to map back to the galaxy that produced it, or it has no ground truth and is useless for training. That mapping is a position match between grizli’s segmentation and the simulation’s per-image truth index.

specsr_roman.extraction.catalog.ab_h158(f_cat)[source]

OU2024 catalogue H158 flux -> AB magnitude.

The images carry counts = f_cat * 10^(0.4 * ZPTMAG), which fixes the anchor at 14.96. Use this rather than the truth index’s own mag column, which is instrumental (-2.5 log10(flux) + const) and only good for ranking.

Return type:

ndarray

specsr_roman.extraction.catalog.load_truth_index(path, zptmag=16.8009)[source]

Per-image truth index -> (ids, ra, dec, x, y, mag) for galaxies only.

Columns are object_id ra dec x y realized_flux flux mag obj_type. Stars and transients are dropped: they have no SED entry in the galaxy catalogue, so they can shape the contamination scene but never be targets.

Parameters:
specsr_roman.extraction.catalog.detect_and_relabel(flt, truth, match_radius=3.0, threshold_sigma=2.5, npixels=8, verbose=True)[source]

Detect sources and relabel the segmentation with compact ids.

Returns (compact_ids, object_ids, mags, matched), all indexed by detection.

Two things are load-bearing here.

Compact ids. grizli keeps the segmentation map as float32, which cannot represent OU2024’s 13-digit object_id exactly — ids silently collide or shift. The segmentation therefore carries a small sequential id (detection index + 2) and object_ids maps back. -1 marks a detection with no truth counterpart; those still shape the contamination scene, they just cannot be targets.

Detection runs on padded arrays. grizli pads the frame so edge sources’ beams stay in-array, so truth coordinates must be shifted by the pad before matching or every match is wrong by a few hundred pixels.

Parameters:

Ground-truth SEDs from the OU2024 (Diffsky) catalogue.

These are the super-resolution targets, so how they are read matters as much as how the spectra are extracted. Two properties of the format decide everything downstream:

  • the wavelength grid is adaptive — sub-Angstrom bins at emission lines, very coarse elsewhere. Point interpolation onto a uniform grid loses line flux, so every rebin here is flux-conserving;

  • the absolute flux scale is internal to the simulation and meaningless. Only the shape is used: grizli renormalises to the direct-image counts, and the training dataset normalises per spectrum.

class specsr_roman.extraction.seds.SEDLibrary(path)[source]

Lazily opened handle on a galaxy_sed_<healpix>.hdf5 file.

The file is ~14 GB per healpix and is read thousands of times per SCA, so it is opened once and the rest-frame wavelength grid cached. Layout is skyCatalogs’: galaxy/<gid // 100000>/<gid> holding (3 components, n_wave), with meta/wave_list the rest-frame wavelengths in Angstrom.

Parameters:

path (str)

observed(gid, z, dlam=5.0)[source]

(observed wavelength [A], f_lambda) on a uniform dlam grid.

The three Diffsky components are summed, redshifted, and rebinned flux-conservingly so the line spikes survive every later interpolation. Raises KeyError for a galaxy with no SED entry.

Parameters:
grizli_spectrum(gid, z)[source]

([wave, flux_normalised], (wave, flux_raw)) for a grizli dispersal.

Returns None if the galaxy has no usable flux in the normalising band. The second element is the un-normalised SED, which is what becomes the training target.

Parameters:

Dispersing a scene and adding Roman grism noise.

specsr_roman.extraction.simulate.disperse_scene(flt, compact_ids, object_ids, mags, redshifts, seds, scene_indices, verbose=True)[source]

Disperse every listed source with its true SED into flt.model.

Returns {index: grizli spectrum} for the sources that succeeded, so a later pass can re-disperse a target on its own without re-reading the SED file.

size=85, compute_size=False is not tunable in practice. The Roman trace sits up to ~66 px from the source on detector 1 (and ~162 px on detector 4), which overflows grizli’s default adaptive cutout and silently truncates the beam.

Parameters:

verbose (bool)

specsr_roman.extraction.simulate.add_grism_noise(scene, exptime=301.0, background=0.57, read_noise=8.5, seed=0)[source]

Noiseless scene (e-/s) -> (noisy, error), both in e-/s.

Poisson from source plus zodiacal background, plus read noise, all expressed as a rate variance so the arrays stay in the units grizli’s optimal extraction expects. The Gaussian approximation to the Poisson term is safe here: even a faint HLSS source accumulates enough electrons over 301 s that the distribution is near-normal, and the background alone contributes ~170 e-.

Parameters:

Contamination-subtracted optimal extraction of a single target.

specsr_roman.extraction.extract.extract_target(flt, scene, noisy, err2d, compact_id, mag, spectrum_1d, min_finite=100)[source]

Extract one target, subtracting everything else in the scene.

Returns (wave, flam, flam_err) on the native grism sampling.

Contamination is handled exactly, not modelled approximately: the target is re-dispersed alone with its own true SED, and that 2D model is subtracted from the full scene. Whatever remains inside the beam is genuinely other galaxies’ light — which is the whole point of simulating a slitless survey rather than isolated sources.

Flux calibration comes from a second, flat-f_lambda pass through the same beam: dividing the extracted counts by the extracted flat model converts to f_lambda while cancelling the trace, the sensitivity curve and the optimal-extraction profile in one step. Pixels where the flat model falls below 5% of its peak are set to NaN — the band edges, where that division is unstable.

Parameters:
exception specsr_roman.extraction.extract.ExtractionFailure[source]

A target could not be extracted usefully. Skip it and move on.

Configuration and constants

Typed configuration for the three training stages.

Every hyperparameter of the published models lives in configs/*.yaml and is loaded into one of the dataclasses below. Two reasons this is not just argparse:

  • a run is reproducible from a file you can diff, rather than from a shell history line;

  • the defaults here are the canonical chain. Constructing SR1Config() with no arguments gives the settings that produced the published SR1 checkpoint, so “what was this trained with” has a readable answer.

Values carried over from the JWST sweep are marked; they were not re-swept on Roman, and re-sweeping them is open work.

class specsr_roman.config.SR1Config(data='data/dataset/ou2024_h10307_dataset.npz', augment=False, min_strong_lines=0, hidden_dim=120, num_res_blocks=16, dropout=0.023538492919758583, in_channels=2, epochs=200, batch_size=32, lr=8.23706977169561e-05, weight_decay=4.953557427559904e-05, grad_clip=0.5, use_var_clamp=True, var_clamp_min=0.1, var_clamp_max=30.0, logvar_reg=1.5590946894260903e-06, mask_smooth_k=161, mask_thresh_mad=8.0, mask_dilate=11, mask_min_width=7, lam_d1=0.11, lam_d2=0.01, gate_min_frac=0.01282614837763313, gate_temp=0.2962018702777486, score_w_line=0.50334848260373, score_w_recon=0.2, sharp_wd_start_epoch=25, sharp_wd_rate=0.008, sharp_wd_floor=0.15, lam_lineflux=1.0, lineflux_snr0=2.0, lam_hallu=0.5, hallu_snr0=1.0, ema_alpha=0.9, init_checkpoint=None, out_prefix='sr1_roman', out_dir='runs/sr1', run_name=None, seed=42, num_workers=4, progress=False, wandb_project='roman-spectral-superresolution', wandb_mode=None, push_to_hub=False, hub_repo=None, export_predictions=True)[source]

SR1: coarse super-resolution.

Architecture and optimiser values are the JWST sweep optima, deliberately unchanged so the two instruments’ models stay comparable. Everything under “recoverability” is Roman-specific and was tuned here.

Parameters:
class specsr_roman.config.ZHeadConfig(data='data/dataset/ou2024_h10307_dataset.npz', sr1_ckpt='sr1_ou2024_v6', min_strong_lines=0, arch='clf', hidden_dim=64, num_blocks=4, dropout=0.1, n_heads=8, n_bins=310, z_lo=0.0, z_hi=3.1, label_sigma=0.03, refine_window=8, use_phot=True, phot_tier='medium', phot_mag_err=0.05, phot_eval_mag_err=0.05, epochs=150, batch_size=32, lr=0.0003, weight_decay=1e-05, z_var_floor=1e-06, out_prefix='zhead_roman', out_dir='runs/zhead', run_name=None, seed=42, num_workers=4, wandb_project='roman-spectral-superresolution', wandb_mode=None, push_to_hub=False, hub_repo=None)[source]

ZHead: P(z) over a redshift grid, conditioned on grism + photometry.

Parameters:
class specsr_roman.config.SR2Config(data='data/dataset/ou2024_h10307_dataset.npz', sr1_ckpt='sr1_ou2024_v6', zhead_ckpt='zhead_ou2024_roman_med3_noisy', phot_tier='medium', min_strong_lines=0, augment=False, epochs=150, batch_size=32, lr=0.0001, weight_decay=2e-05, grad_clip=0.5, delta_cap=40.0, sigma_base_um=0.005, z_topk=3, lam_hp_in=3.0, lam_hp_out=0.3, hp_k=51, lam_sparse=0.0, lam_lineflux=1.0, lam_hallu=1.0, hallu_snr0=1.0, lam_presence=0.3, presence_thresh=3.0, label_sigma=0.03, lam_z=0.0, lam_z_warmup=5, zhead_finetune=True, zhead_lr_mult=0.1, out_prefix='sr2_roman', out_dir='runs/sr2', run_name=None, seed=42, num_workers=4, wandb_project='roman-spectral-superresolution', wandb_mode=None, push_to_hub=False, hub_repo=None, smoke=False)[source]

SR2: line-token attention refinement on top of frozen SR1 + ZHead.

Parameters:
specsr_roman.config.load_config(cls, path=None, overrides=None)[source]

Build a config from an optional YAML file plus optional overrides.

Unknown keys are an error rather than a shrug: a typo’d hyperparameter that silently does nothing is the worst possible failure mode for a training run you will not look at again for six hours.

Parameters:
Return type:

T

specsr_roman.config.to_dict(cfg)[source]

Dataclass -> plain dict, for W&B config and checkpoint cards.

Return type:

dict[str, Any]

Checkpoint resolution and loading.

A checkpoint may be named three ways, and every entry point accepts all three:

  • a local path — runs/sr1/sr1_ou2024_v6_best.pth;

  • a bare run name — sr1_ou2024_v6, fetched from the Hugging Face Hub;

  • a fully qualified hub reference — org/repo:run_name.

Weights are downloaded once into the standard Hugging Face cache and reused, so the published pipeline runs from a clean checkout with no manual downloads.

specsr_roman.checkpoints.CANONICAL_CHAIN = {'sr1': 'sr1_ou2024_v6', 'sr2': 'sr2_ou2024_v5_romanonly', 'zhead': 'zhead_ou2024_roman_med3_noisy'}

The three checkpoints that make up the published pipeline. Anything else on the Hub is a superseded generation kept for provenance.

specsr_roman.checkpoints.resolve_checkpoint(spec, repo_id=None)[source]

Checkpoint spec -> a local file path, downloading from the Hub if needed.

Parameters:

  • spec (str)

  • repo_id (str | None)

Return type:

Path

specsr_roman.checkpoints.load_sr1(spec=None, device='cpu', repo_id=None, hidden_dim=120, num_res_blocks=16)[source]

Frozen SR1, ready for inference. Input width is read off the checkpoint.

Parameters:

  • spec (str | None)

  • repo_id (str | None)

  • hidden_dim (int)

  • num_res_blocks (int)

specsr_roman.checkpoints.load_zhead_ckpt(spec=None, device='cpu', repo_id=None)[source]

Whichever ZHead generation the checkpoint holds, ready for inference.

Parameters:

  • spec (str | None)

  • repo_id (str | None)

specsr_roman.checkpoints.load_sr2(spec=None, device='cpu', repo_id=None, in_channels=6, line_rest_um=None, wave_hi_um=None)[source]

Frozen SR2. Falls back to the packaged line list and HR grid.

Parameters:

  • spec (str | None)

  • repo_id (str | None)

  • in_channels (int)

specsr_roman.checkpoints.push_checkpoint(pth_path, run_name, meta=None, repo_id=None)[source]

Upload a trained checkpoint plus a provenance card to the Hub.

The card is not optional decoration. Local disk on a shared machine is not a backup, and a .pth with no record of the dataset, the upstream checkpoints, the config and the W&B run is an unreproducible artefact within a week.

Parameters:
Return type:

str

Fixed wavelength grids, instrument constants, and photometric band maps.

Every spectrum in the dataset lives on one of two shared grids, so a model is fully convolutional over a fixed-length axis and any two rows are directly comparable. The grids mirror the JWST/JADES convention of the companion project (specsr) — same npz keys, same “LR interpolated onto the HR grid at load time” contract — which is what made the cross-instrument warm-start experiments possible.

specsr_roman.grids.resolve_phot_tier(spec)[source]

"medium" or an explicit "8,9,11" -> band indices.

None means “use every band the dataset file carries” and is a loader convenience — OU2024 stores 14 columns whatever a model consumes. It is not a model configuration: every training config names a tier, and MAX_PHOT_BANDS caps what a model can be handed.

Parameters:

spec (str | None)

Return type:

tuple[int, …] | None

Rest-frame spectral features used across the pipeline.

Three line sets appear here, and they are not interchangeable:

LINE_LIST_REST_AA

The full 98-feature list that SR2’s line-token transformer attends over. Breadth matters more than in-band coverage: a token whose line falls outside the grism band is gated off by in_range in the line branch, so an over-complete list costs a little compute and buys robustness across the redshift range.

SR1_LINES_AA

The ten redshift-carrying lines SR1 supervises directly (line-flux L1 and the anti-hallucination penalty), and the lines the dataset computes per-row integrated S/N labels for. Order is load-bearing — the line_snr column of specsr_roman.data.RomanFixedGridDataset follows this tuple, and the losses index into it positionally.

STRONG_LINES_AA

The four lines that decide whether a redshift is identifiable at all. Two or more in band means the redshift is line-pair constrained; one is alias-degenerate (Ha/[OIII]/[OII] all look alike on their own); zero (z < 0.52 for the Roman band) means unconstrained. Used by the min_strong_lines split filter.

specsr_roman.lines.angstrom_to_micron(x)[source]

Angstrom -> micron, as float32 (the dtype every model buffer uses).

Return type:

ndarray

specsr_roman.lines.count_strong_lines(z, lam_min_aa, lam_max_aa)[source]

Number of STRONG_LINES_AA inside [lam_min_aa, lam_max_aa] at z.

Parameters:
Return type:

ndarray

Where things live.

Resolution order for every data root is: explicit argument, then environment variable, then a sensible default relative to the current working directory. Nothing here reaches outside the project, and nothing hard-codes a machine.

specsr_roman.paths.data_root()[source]

Root for downloaded and prepared simulation products (SPECSR_ROMAN_DATA).

Return type:

Path

specsr_roman.paths.dataset_dir()[source]

Where built training datasets land (SPECSR_ROMAN_DATASETS).

Return type:

Path

specsr_roman.paths.runs_dir()[source]

Training outputs: checkpoints, predictions, split records (SPECSR_ROMAN_RUNS).

Return type:

Path

specsr_roman.paths.cache_dir()[source]

Scratch for extraction intermediates (SPECSR_ROMAN_CACHE).

Return type:

Path

specsr_roman.paths.grizli_conf_dir()[source]

Local grizli CONF tree. grizli reads $GRIZLI, so honour it first.

Return type:

Path