specsr.models.zhead

ZHead — redshift inference from a super-resolved spectrum (stage 2 of 3).

Consumes the SR1 output and predicts a continuous redshift with an associated uncertainty. In the full pipeline the predicted redshift conditions the SR2 line prior, so its role is structural, not merely diagnostic: it tells SR2 where in wavelength the emission lines should be.

The same architecture is also trained separately on low-resolution, super- resolved and high-resolution inputs. Comparing those three runs isolates how much redshift-relevant information each representation carries, holding architecture, loss and split fixed.

Functions

heteroscedastic_nll(mu, log_var, y[, var_floor])

Gaussian negative log-likelihood with a predicted per-sample variance.

redshift_pdf_loss(logits, z_true_n, z_grid_n)

Cross-entropy of a predicted redshift PDF against a soft target.

Classes

ZHead1D(*args, **kwargs)

Continuous redshift head with uncertainty-aware pooling.
class specsr.models.zhead.ZHead1D(*args, **kwargs)[source]

Bases: Module

Continuous redshift head with uncertainty-aware pooling.

Parameters:

  • in_channels (int) – 1 to consume flux alone, 2 to additionally consume the SR1 predictive log-sigma. With 2 channels the network routes flux and uncertainty through separate paths and pools by confidence.

  • hidden_dim (int) – Width of the flux path. The uncertainty path uses hidden_dim // 2, since it carries less information.

  • num_blocks (int) – Number of convolutional blocks per path.

  • dropout (float) – Dropout probability inside each block.

  • coord_channel (bool) – Append a normalised wavelength coordinate (linspace(-1, 1, L)) as an extra input channel. v2 only.

  • dilation_growth (int) – Dilation of block i is dilation_growth ** i. At 1 (default) this is the v1 architecture; at 2 with six blocks the receptive field grows from 37 to ~380 pixels. v2 only.

  • pooling (str) – "mean" selects the original (v1) architecture below. "attn" selects the v2 architecture: a single trunk over all input channels followed by attention pooling over wavelength.

  • head (str)

  • n_z_bins (int)

  • z_grid_min_n (float | None)

  • z_grid_max_n (float | None)

  • soft_argmax_half (int)

Returns:

Both shaped (B,).

Return type:

mu, log_var

Notes

v1 (``pooling=”mean”``) is redshift-blind by construction and is kept only so historical checkpoints load. Its trunk is translation-equivariant and its global mean pool is translation-invariant, so a spectrum whose lines are shifted – which is what a different redshift is on a fixed log-wavelength grid – produces near-identical pooled features. The three-arm comparison run on it produced the impossible ordering hires < lowres, because the only signal it could use was shift-invariant statistics (line strengths, continuum shape), which the prism carries at higher per-pixel S/N.

v2 (pooling="attn") fixes this three ways: a coordinate channel tags every feature with its position on the grid, dilated convolutions widen the receptive field so multi-line patterns are co-detected, and attention pooling forms a position-weighted summary instead of a position-destroying average. When the uncertainty channel is present it is treated as one more input channel of the same trunk, so all comparison arms share one structure and differ only in channel count.

In v1, when the uncertainty channel is present, the pooling over wavelength is weighted by predicted confidence rather than being a plain mean: confidence is derived as inverse sigma, smoothed by a wide convolution, then normalised to sum to one across the spectral axis.

property bounded_mean: bool

Whether forward’s first output still needs decode_mean.

The Gaussian head emits an unbounded scalar that the transform squashes onto the observed redshift range. The classification head’s estimate is already in normalised redshift units – it comes from a grid that cannot leave the range by construction – so squashing it again would be wrong. Callers branch on this rather than on the head name.

forward(x)[source]

(mu, log_var), both (B,).

For the classification head these are derived from the predicted PDF, so every existing consumer keeps working unchanged; train against logits() instead, which is what the objective needs.

Parameters:

x (torch.Tensor)

Return type:

tuple[torch.Tensor, torch.Tensor]

logits(x)[source]

Unnormalised log-PDF over the redshift grid, (B, n_z_bins).

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

moments(logits)[source]

PDF -> (mu, log_var) in normalised redshift units.

The point estimate is a local expectation in a window around the modal bin, not the global mean. That distinction is the whole reason for a classification head: redshift PDFs are genuinely multimodal (an emission line misidentified as a different transition puts a second peak elsewhere), and the global mean of a bimodal PDF lands in the empty valley between the two modes – which is precisely the catastrophic-outlier failure this head exists to remove. Taking a soft-argmax keeps the estimate on a mode while staying differentiable, so the redshift coupling in the SR2 loss can still backpropagate.

log_var is computed from the full PDF, so it stays large when the head is genuinely torn between two redshifts – which is the signal downstream stages need in order to distrust it.

Parameters:

logits (torch.Tensor)

Return type:

tuple[torch.Tensor, torch.Tensor]

specsr.models.zhead.heteroscedastic_nll(mu, log_var, y, var_floor=1e-06)[source]

Gaussian negative log-likelihood with a predicted per-sample variance.

All arguments are shaped (B,). The variance floor prevents the loss from diverging when the model becomes overconfident on easy examples.

Parameters:
Return type:

torch.Tensor

specsr.models.zhead.redshift_pdf_loss(logits, z_true_n, z_grid_n, target_sigma_bins=1.5)[source]

Cross-entropy of a predicted redshift PDF against a soft target.

logits is (B, n_bins), z_true_n is (B,) in normalised redshift units, z_grid_n is the bin centres.

The target is a narrow Gaussian on the grid rather than a one-hot bin. With ~10^3 bins the true redshift almost never sits exactly on a centre, and one-hot targets make the objective discontinuous in the true value while telling the model that the neighbouring bin – which may be a small fraction of a resolution element away – is exactly as wrong as the far end of the grid. A soft target keeps the loss smooth and encodes that near misses are near.

Unlike a Gaussian likelihood this objective is indifferent to how many modes the posterior has, which is the point: it can represent “either z=2.1 or z=5.3” honestly instead of splitting the difference.

Parameters:
Return type:

torch.Tensor