specsr.models.sr2

SR2 — physics-informed residual refiner (stage 3 of 3).

Predicts a correction delta to the SR1 output, as the sum of two branches:

Line branch. For each of K catalogued rest-frame features, the observed position is computed from the redshift predicted by ZHead. A window of the input is extracted around each position and encoded into a token; multi-head self-attention is applied across line tokens, so the model can exploit physical relationships between lines (fixed doublet ratios, Balmer decrements, lines that co-occur in a given ionisation state). Each token then emits Gaussian parameters — amplitude, width, sub-pixel offset — gated by a learned presence probability, and the Gaussians are rendered back onto the wavelength grid.

CNN branch. A plain residual CNN handles smooth continuum corrections that the parametric line model cannot express.

The parametric line branch is what makes this “physics-informed”: the model cannot place a sharp feature at an arbitrary wavelength, only at a catalogued line position implied by the redshift. That constraint is the main defence against the failure mode of end-to-end super-resolution models, which tend to either over-smooth real features or hallucinate sharp ones.

Warning

Attribute names here are load-bearing. Published checkpoints key on line_rest_um and wave_hi_um as registered buffers; an earlier iteration used line_rest_um_buf/wave_hi_um_buf and its weights will not load against this class without remapping.

Functions

build_line_mask(wave_hi_um, zhat, line_rest_um)

Per-sample soft mask marking where emission lines are expected.

build_sr2_input(x_low, sr1_mean, ...[, ...])

Assemble the channel stack SR2 consumes.

constrain_delta(delta, cap)

Soft-limit the residual magnitude to +/- cap.

gaussian_line_mask(wave_um, centers_um, sigma_um)

Union of Gaussians centred on centers_um, shaped (1, 1, L).

sr2_input_channels([use_sr1_sigma, ...])

Number of channels build_sr2_input() will produce.

Classes

SR2Attention(*args, **kwargs)

Two-branch residual refiner: line-token attention plus a residual CNN.
class specsr.models.sr2.SR2Attention(*args, **kwargs)[source]

Bases: Module

Two-branch residual refiner: line-token attention plus a residual CNN.

Parameters:

  • in_channels (int) – Channels of the input stack handed to SR2 (typically the SR1 mean, its log-sigma, and the interpolated low-resolution input).

  • line_rest_um – (K,) rest-frame line wavelengths in microns. One attention token per entry.

  • wave_hi_um – (L,) observed-frame wavelength grid in microns.

  • line_dim (int) – Token embedding width for the line branch.

  • num_attn_heads (int) – Transformer encoder shape for cross-line attention.

  • num_attn_layers (int) – Transformer encoder shape for cross-line attention.

  • window_half (int) – Half-width, in pixels, of the window extracted around each line.

  • cnn_dim (int) – Width and depth of the continuum CNN branch.

  • num_cnn_blocks (int) – Width and depth of the continuum CNN branch.

  • dropout (float) – Dropout probability in both branches.

  • cnn_scale (float)

Returns:

  • In training mode, (delta, logvar, presence, presence_logit); in eval

  • mode, (delta, logvar). presence is exposed during training so the

  • loss can supervise it against the lines the reference spectrum actually

  • shows — without that signal the model turns every line on weakly rather than

  • committing to the ones that are really there.

Notes

Initialisation is deliberately asymmetric between the branches, and it is worth being precise about because it is not a clean identity:

  • The CNN branch output head is fully zero-initialised (weight and bias), so it contributes exactly zero at step zero.

  • The line branch is not. amp_head has zero bias but random weights (std=0.01), and the presence gate starts at sigmoid(-2) ~ 0.12, so the branch injects small random Gaussians at every catalogued line position from the first forward pass. Measured on a randomly initialised model, max|delta| ~ 0.03 in normalised flux units.

So SR2 begins as a near-identity on top of SR1, not an exact one. The perturbation is small relative to the normalised flux scale, but it is the amplitude head — the component that sets emission-line flux — that carries it.

Note

Flux conservation is the known weak point of this architecture: it places lines at the right wavelengths but does not necessarily preserve their integrated flux. Anything touching amp_head, the presence gate, or the residual cap should be evaluated against an SR-versus-HR line flux comparison (scripts/flux_conservation.py), not against positional accuracy or S/N alone.

One mechanism, found on an earlier checkpoint, is worth stating plainly, because it is a property of this amp * presence factorisation rather than a tuning accident. Under a blanket sparsity penalty the presence gate settled at ~0.002 for real and absent lines alike – a discrimination ratio of 0.95, i.e. none – so it stopped acting as a selector and became a constant multiplier on every amplitude. The line branch then supplied 0.17% of the flux the reference required, and every gain SR2 showed on lines was coming from the CNN branch instead. The amplitudes themselves were fine: amp_head was emitting ~7 in normalised units where the reference wanted ~8.9. Only the gate was wrong. Because the two are multiplied, any pressure on presence is pressure on flux, so presence must be supervised rather than merely penalised.

line_positions(zhat)[source]

Public view of _line_positions(), for losses that need to line their windows up with where this model placed its Gaussians.

Parameters:

zhat (torch.Tensor)

Return type:

tuple[torch.Tensor, torch.Tensor]

forward(x, zhat)[source]
Parameters:
specsr.models.sr2.constrain_delta(delta, cap)[source]

Soft-limit the residual magnitude to +/- cap.

A tanh squash rather than a hard clamp, so gradients keep flowing when the correction saturates. Capping matters because SR2 is free to add flux anywhere; without a bound it can manufacture arbitrarily large features.

Parameters:
Return type:

torch.Tensor

specsr.models.sr2.gaussian_line_mask(wave_um, centers_um, sigma_um)[source]

Union of Gaussians centred on centers_um, shaped (1, 1, L).

sigma_um is a scalar, or a per-line tensor broadcastable to centers_um – the latter is what lets a mask widen line-by-line with the redshift uncertainty, since a given redshift error displaces a red line further in wavelength than a blue one.

Parameters:
Return type:

torch.Tensor

specsr.models.sr2.build_line_mask(wave_hi_um, zhat, line_rest_um, sigma_base_um=0.005, z_sigma=None, sigma_max_um=0.05)[source]

Per-sample soft mask marking where emission lines are expected.

Used both as an SR2 input channel and by the loss to weight line regions differently from continuum.

z_sigma widens the mask by the redshift uncertainty. A line of rest wavelength lam0 sits at lam0 * (1 + z), so an uncertainty sigma_z displaces it by lam0 * sigma_z; that is combined with the base width in quadrature. Without it the mask is a fixed 0.005 um – a redshift precision of dz ~ 0.01 – and is simply wrong whenever the head is less certain than that, pointing SR2’s attention confidently at a stretch of continuum. Widening makes an unsure prediction produce a vague mask instead of a confident false one.

sigma_max_um caps the widening. The same mask weights the in-line term of the SR2 loss, so an uncapped mask on a hopeless redshift would approach all-ones and quietly turn that term into a plain reconstruction loss.

Parameters:
Return type:

torch.Tensor

specsr.models.sr2.build_sr2_input(x_low, sr1_mean, sr1_log_sigma, zhat, wave_hi_um, line_rest_um, *, use_sr1_sigma=True, use_line_mask=True, use_zhat_channel=True, use_zsigma_channel=False, sigma_base_um=0.005, z_sigma=None, sigma_max_um=0.05)[source]

Assemble the channel stack SR2 consumes.

Channel order is fixed and load-bearing — a checkpoint trained on one order produces silent nonsense if fed another:

  1. x_low — the interpolated low-resolution input

  2. sr1_mean — the SR1 reconstruction

  3. sigma — exp(sr1_log_sigma), if use_sr1_sigma

  4. line_mask— soft mask of expected line positions, if use_line_mask

  5. zhat — the redshift broadcast along wavelength, if use_zhat_channel

  6. z_sigma — its uncertainty, likewise broadcast, if use_zsigma_channel

z_sigma is appended last so that enabling it does not renumber the existing channels. It still changes the input width, so a checkpoint trained without it cannot be warm-started with it on. Widening the line mask via z_sigma needs no such retraining – the channel count is unchanged – which is why the two are separate flags.

Important

Channel 2 is the linear predictive sigma, not the log-sigma. SR1 emits log_var; callers pass log_sigma = 0.5 * log_var and this function exponentiates it. Passing the log by mistake is dimensionally plausible and trains to a worse but non-obviously-broken model.

Giving SR2 both the line mask and the redshift explicitly, rather than letting it infer them, is part of what makes the stage physics-informed: the conditioning is supplied rather than learned.

This function exists so the stack is defined exactly once. It was previously open-coded in the SR2 training script, the inference script and two redshift head variants, where the copies could drift apart silently.

Parameters:
Return type:

torch.Tensor

specsr.models.sr2.sr2_input_channels(use_sr1_sigma=True, use_line_mask=True, use_zhat_channel=True, use_zsigma_channel=False)[source]

Number of channels build_sr2_input() will produce.

Parameters:

  • use_sr1_sigma (bool)

  • use_line_mask (bool)

  • use_zhat_channel (bool)

  • use_zsigma_channel (bool)

Return type:

int