specsr — spectral super-resolution for JWST/NIRSpec galaxy spectra

Weights for the pipeline described in:

Haghjoo, A., Hemmati, S., Mobasher, B., et al. Learning to See Sharper: A Physics-Informed Artificial Intelligence Framework for Super-Resolving Galaxy Spectra. ApJ, under revision (AAS75211). arXiv:2603.18357

main serves the current chain (v4, 2026-08-10). Three earlier revisions are kept as archival tags for reproducibility and must not be used for science — see Revisions.

from specsr.inference.pipeline import SpecSRPipeline

pipe = SpecSRPipeline.from_pretrained()      # downloads ~18 MB
result = pipe(flux_low)                      # on the package's log constant-R grid
result.sr2, result.z, result.z_sigma

What the pipeline does

Every number on this diagram — channel widths, block counts, the 98 line tokens, the 161-pixel window, the 1,024 redshift bins, the per-stage parameter totals — is read off the checkpoints in this repository at draw time, and the three inset panels are real predictions for one held-out galaxy at z = 2.198. Regenerate with scripts/make_architecture_figure.py in the code repository.

Three stages, trained in order:

Stage	Role
SR1	1D ResNet-CNN backbone, 16 blocks x 120 ch: prism (R100) → grating-like (R1000) reconstruction, with a per-pixel log-variance
ZHead	Redshift inference from SR1's mean and log-variance; softmax PDF over 1,024 bins spanning 0 <= z <= 15, point estimate by soft-argmax on the mode
SR2	Physics-informed residual refiner: attention over 98 emission-line tokens, each a parametric Gaussian gated by a supervised presence probability, plus a CNN continuum branch

SR1 is frozen while SR2 trains. ZHead is not quite frozen — its attention-pooling layer and its 1,024-way output layer train at 0.1x the SR2 learning rate (zhead_unfreeze_last_n: 2 counts parameterised leaf modules, not blocks: 133,249 of 710,145 weights; the conv trunk is untouched).

That adaptation makes the redshift head worse, and the adapted head is not what you download. Measured on the 572 held-out galaxies, the outlier rate goes 10.84% frozen -> 14.69% after one epoch of SR2 training -> 17.83% by epoch 30, consistently across three runs. The SR2 trainer then persists only sr2_state_dict, so zhead/best_zhead.pth here is the head as it stood before that fine-tuning — the 10.84% one.

Two consequences worth being explicit about. Every number on this card was measured through exactly the chain you download, so the reported performance is what you get. And the direction of the mismatch is favourable: SR2 trained against a ~14% head and runs against a ~10.8% one, so at inference it receives better redshifts than it was trained with. Future revisions will simply freeze the head outright.

There is no teacher forcing anywhere in the chain: SR2 only ever sees redshifts ZHead predicted, never the catalogue value.

Honest performance summary

All numbers on 572 held-out galaxies, split by parent galaxy so no augmented sibling crosses the split. JADES DR4, log constant-R grid (R_grid = 4000).

Emission-line flux (SR2 vs HR reference)

The SR/HR flux ratio is a strong function of line SNR — a single number is misleading. Median ratio by HR SNR cut: 0.24 / 0.51 / 0.61 / 0.67 at 5 / 10 / 20 / 50. Below SNR ~20 the HR line flux is itself mostly noise.

HR SNR cut	n	median SR/HR	16-84 scatter	within 20%	within 50%
>= 5	5368	0.244	[-0.01, +0.88]	12.5%	30.6%
>= 10	2423	0.512	[+0.05, +1.01]	20.3%	46.7%
>= 20	1203	0.609	[+0.09, +1.02]	25.7%	54.2%
>= 50	350	0.671	[+0.08, +1.03]	34.6%	62.3%

Read the scatter, not just the median: only ~35% of even the brightest lines land within 20% of the true flux. Usable in aggregate for bright lines; not yet reliable line by line. Blended systems (Halpha+[N II], the [S II] doublet) carry larger uncertainty than isolated lines — 11 of the 98 catalogued lines have no line-free continuum window within +/-1500 km/s.

Deblending — the headline claim, measured

The [O III] 4959/5007 doublet is the clean test of whether super-resolution makes blended structure usable, because its ratio is fixed by atomic physics at 2.98 — no ground truth is needed and any departure is measurement error. On 318 held-out galaxies with the reference doublet detected at SNR >= 20, every input gets the same classical fit (linear continuum + three Gaussians at known positions, shared width, non-negative amplitudes):

input	within 20% of 2.98	16-84 range	fitted sigma
Low-res prism (raw input)	39.3%	1.47-5.88	1848 km/s
SR1 output	14.3%	0.40-10.10	423 km/s
SR2 output	32.4%	1.58-6.44	365 km/s
High-res grating (reference)	85.8%	2.65-3.25	168 km/s

Neither SR1 nor SR2 recovers the doublet ratio better than directly fitting the prism, and SR1 is substantially worse than its own input. Paired on the same galaxies, above SNR 50 the prism beats SR2 39:17 (p=0.005); the 39:28 lead SR2 holds at SNR 20-50 is not significant (p=0.22). There is no SNR regime where SR2 significantly wins, and extending down to SNR 3 does not find one — below SNR ~10 the ratio is recovered poorly from every input including the grating itself.

The interpretation: the prism is photon-limited, so its accuracy improves as lines brighten, while SR2 is limited by its own reconstruction error and stays flat. Super-resolution places lines at the correct wavelengths and separates them visually — which is what drives the SNR gains — but does not thereby make the amplitude ratio of two blended components more faithful. One quantity does improve: SR2 fails to detect 4959 at all in 7.9% of galaxies against the prism's 15.1%, so the reconstruction is more robust even where it is not more accurate.

Reproduce with scripts/doublet_deblending.py in the code repository.

Redshift — and a limitation to be aware of

One head architecture, one split, one seed; only the input representation differs. Outliers are |dz|/(1+z) > 0.15.

head input	outliers	med |dz|/(1+z)	sigma_NMAD
Low-res prism (raw input)	6.12%	0.00151	0.00223
SR1 output	10.84%	0.00177	0.00269
High-res grating (reference)	11.01%	0.00083	0.00125
SR2 output	12.06%	0.00145	0.00223

Super-resolution does not currently improve redshift recovery over the raw prism input. We state this plainly because the repository's own numbers say so. The SR2 arm, retrained against this revision's SR2 and added on 2026-08-10, completes the picture rather than changing it: it is the best of the four on median precision after the grating, and the worst of the four on catastrophic outliers, roughly 2x the prism's rate.

Read the SR2 arm with a caveat the other three do not need. SR2 renders its line Gaussians at the redshift the SR1 head inferred, so a redshift estimated from an SR2 spectrum is not an independent measurement — it substantially recovers a value the pipeline already committed to. It agrees with that conditioning redshift to <1% for 86% of galaxies; its outlier rate is 85.5% where the SR1 head was wrong and 3.1% where it was right, which accounts for the 12.06% in full. Of the 62 galaxies the SR1 head misidentifies it recovers 9, against the grating arm's 38 on the same objects. So SR2's outliers are inherited identifications rather than the grating's faint-line detection failures — only 41% of the two arms' outliers overlap. Restricted to correctly identified galaxies the refinement does what it is meant to: median |dz|/(1+z) improves from SR1's 0.00149 to 0.00128, ahead of the prism's 0.00136. Sharpening helps localisation and structurally cannot repair a misidentified line system.

An earlier sr2 arm exists in our run directory and is not published here: it was trained against a superseded SR2 and its numbers do not describe this chain.

The pattern is interpretable rather than mysterious: the prism and the grating trade SNR against resolution at fixed exposure — the prism has 4.4x higher median per-pixel SNR (4.13 vs 0.94), the grating ~4x the resolving power. Precision follows resolution (the grating is 2x better on median |dz|/(1+z)); catastrophic outliers follow SNR (the prism is ~2x better). Super-resolution's goal is to obtain both, and on this measurement it does not yet.

Use the prism-trained redshift head, not the SR1 head, if redshift is what you need. Both are in this repo.

Contents

Path	Stage
sr1/best_sr1.pth	SR1 backbone (selected)
sr1/config_logR.yaml	SR1 architecture config — required to rebuild SR1
zhead/best_zhead.pth	ZHead on SR1 output (softmax PDF head)
sr2/best_sr2.pth	SR2 residual refiner
zhead/best_zhead_lowres.pth	ZHead on raw prism input — comparison arm
zhead/best_zhead_hires.pth	ZHead on true grating spectra — comparison arm
zhead/best_zhead_sr2.pth	ZHead on SR2 output — comparison arm
architecture.png	the pipeline diagram above

The three comparison arms share one architecture, split and seed and differ only in what they are fed; they exist to measure how much redshift information each representation carries, and are not drop-in replacements for zhead/best_zhead.pth, which reads SR1 output and is the one the pipeline uses.

Known limitations

• Redshift: see above. SR1 does not beat the raw prism input.
• Line flux: large per-line scatter (see above). The improvement over previous revisions is in the median bias; the dispersion is essentially unchanged.
• Deblending does not beat direct fitting of the prism — see above. This is the most important limitation for anyone intending to use the reconstructions for line-ratio diagnostics.
• The train/held-out flux gap is largely the redshift head, not SR2. The training flux ratio reaches ~0.93 against ~0.54 held out. This was previously attributed to SR2 memorising amplitudes across the ~21x augmented copies; measurement shows the frozen ZHead is 2.6x more accurate on galaxies it trained on (median |dv| 205 vs 532 km/s) and 21x less catastrophic (0.9% vs 18.9% beyond 2500 km/s). SR2 places its line windows at the predicted redshift, so it is trained in an easy regime and evaluated in a hard one. Restricting flux supervision to original rows would not fix this: augmented and original training rows sit at 230 and 205 km/s respectively.
• ZHead sigma_z is optimistically narrow. The head's point estimate improves while its validation NLL rises, so treat z_sigma as a lower bound.

Revisions

Revision	Status
main (v4, 2026-08-10)	current — use this
v3-20260731	archival. Same SR1 and ZHead as v4; its SR2 was trained with a line mask whose width filter was sized for a retired wavelength grid and which missed 54% of the lines it should have supervised. Superseded.
v2-presencefix-20260726	archival. Trained with unnormalised HR errors, and its presence gate is functionally collapsed (see below). Not for science.
v1-submission	archival. Reproduces the originally submitted manuscript. Trained on a leaky split: 99.2% of galaxies had rows on both sides. Not for science.

The archival tags exist so the published manuscript's numbers remain reproducible while the paper is under review. Neither is reachable by default — DEFAULT_REVISION = "main".

What changed in v4

• SR2 retrained with a corrected presence/flux line mask. The HR-derived mask that labels which catalogued lines are really present — and which lines the flux term sees at all — was built with kernel sizes carried over from a retired 2,500-point linear grid. Measured against integrated HR line SNR, it flagged only 46% of the lines at SNR >= 20; the rest were labelled absent for the presence gate and dropped from the flux term. The cause was the minimum-width filter, not the detection threshold: an R=1000 line on the R=4000 grid clears the threshold over only ~5-6 consecutive samples, and the filter demanded 7.
• What that fixed, honestly. Line flux moved from the generic CNN branch to the physics-informed line branch (0.643 -> 0.676 and 0.343 -> 0.325, both significant paired), and MSE against the reference improved. Total flux conservation did not change (0.618 -> 0.609, paired difference consistent with zero, 48.3% of lines improved). It is a correctness fix with a measured architectural benefit, not a flux improvement.
• The deblending measurement above is new in v4, and it is a negative result.

What changed in v3

• The presence gate is supervised, not penalised. In earlier revisions a blanket sparsity penalty drove it to a constant ~0.002 — identical on real and absent lines, a discrimination ratio of 0.95, i.e. carrying no information. Since the model forms amp * presence, that constant multiplied every line amplitude and the line branch contributed 0.17% of the required flux; every apparent line gain was coming from the generic CNN branch. Presence is now trained by class-balanced BCE against the lines the reference actually shows. Discrimination 16.1x -> 841x.
• An explicit integrated-line-flux term, measured against a line-free continuum so blended systems are not silently mis-measured.
• Retrained on JADES DR4 with a log constant-R grid (the previous linspace grid downsampled the HR targets 2.5x below 3.2 um) and a group-wise 80/20 split.

License

MIT, matching the source repository.