specsr.training.losses

Loss functions for the SR1 and SR2 stages.

Both stages optimise a heteroscedastic Gaussian likelihood against the medium-resolution reference, plus terms that specifically target line structure. The likelihood alone is not enough: it is dominated by the many continuum pixels, and a model that produces a smooth, well-calibrated reconstruction with no sharp lines scores well on it. The extra terms are what make the model commit to narrow features.

A recurring pattern here is gating: extra pressure is applied only where the reference actually shows line structure, and is scaled down for spectra that have essentially none, so featureless objects are not penalised for staying smooth.

Functions

finite_diff(x)

First difference along the last axis, length preserved.

finite_diff2(x)

Second difference along the last axis, length preserved.

gather_line_windows(x, positions, half)

Gather +/-half samples around each line: (B,1,L) -> (B,K,2*half+1).

keep_only_wide(mask[, min_width])

Drop mask features narrower than min_width samples.

line_presence_target(x_high, positions, *[, ...])

Which catalogued lines the reference spectrum actually shows.

line_window_flux(x, positions, *[, ...])

Continuum-subtracted integrated flux per line, (B,K).

make_line_mask_from_smoothed(x_high_raw[, ...])

Locate line regions in the reference spectrum.

masked_mean(x[, valid])

Mean over valid samples only.

robust_mad(x[, dim, eps])

Median absolute deviation, scaled to be a consistent estimator of sigma.

sr1_deblend_loss(mean, log_var, x_high, ...)

Heteroscedastic likelihood plus gated sharpness matching, for SR1.

sr2_loss(*, sr2_mean, sr2_logvar, x_high, ...)

Likelihood, high-pass matching, presence sparsity and a redshift term.
specsr.training.losses.masked_mean(x, valid=None)[source]

Mean over valid samples only.

valid marks wavelengths the reference actually measured. Samples where the detector recorded nothing must contribute nothing to the objective: the original pipeline filled them with the per-spectrum median, which trained the model to emit a flat continuum there and penalised it for placing a real line at a wavelength nobody observed. Excluding is the only honest option, since any fill value is a fabricated target.

Parameters:
Return type:

torch.Tensor

specsr.training.losses.robust_mad(x, dim=-1, eps=1e-08)[source]

Median absolute deviation, scaled to be a consistent estimator of sigma.

Used instead of the standard deviation because emission lines are exactly the kind of large outlier that would inflate it — the scale must describe the noise, not the signal we are trying to detect.

Parameters:
Return type:

torch.Tensor

specsr.training.losses.finite_diff(x)[source]

First difference along the last axis, length preserved.

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

specsr.training.losses.finite_diff2(x)[source]

Second difference along the last axis, length preserved.

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

specsr.training.losses.keep_only_wide(mask, min_width=7)[source]

Drop mask features narrower than min_width samples.

Single-sample spikes are cosmic rays or noise, not resolved emission lines. Requiring a minimum width keeps the line mask from chasing them.

Parameters:
Return type:

torch.Tensor

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

Locate line regions in the reference spectrum.

The continuum is estimated by heavy smoothing and subtracted; what remains is thresholded at thresh_mad robust sigma, filtered to remove narrow spikes, and dilated so the mask covers the line wings rather than only its core.

Deriving the mask from the reference rather than from a line list means it responds to lines that are actually present at the observed strength, and does not mark lines that fall in the wavelength range but are undetected.

Warning

These defaults are sized for the retired 2,500-point linear grid. All four are measured in samples, so they mean different things on the 6,671-point log grid, and min_width in particular is the difference between flagging 46% and 74% of the lines with HR SNR >= 20. Every caller on the log grid must pass its own values – sr1_deblend_loss() takes them from config, and sr2_loss() passes presence_mask_*, which it did not until 2026-08-02. The defaults are kept only so the old-grid behaviour stays reproducible.

Parameters:
Return type:

torch.Tensor

specsr.training.losses.gather_line_windows(x, positions, half)[source]

Gather +/-half samples around each line: (B,1,L) -> (B,K,2*half+1).

Parameters:
Return type:

torch.Tensor

specsr.training.losses.line_window_flux(x, positions, *, core_half=7, sb_lo=11, sb_hi=20, all_positions=None, clean_half=7, min_clean=5)[source]

Continuum-subtracted integrated flux per line, (B,K).

The continuum is the median of two sidebands flanking the core, so the per-spectrum additive normalisation offset cancels and only the common multiplicative scale survives. That is what makes an SR-versus-HR ratio meaningful in normalised units, without de-normalising first.

Summation is in sample units, with no dlambda factor. On this grid dlambda varies across the window, but it is the same at every sample for both spectra being compared, so it cancels from the relative residual the flux loss actually uses. Including it would only rescale each line.

Blended lines. A plain sideband median assumes the sidebands hold only continuum, which is false for a large minority of the catalogue: 11 of the 98 lines have a neighbour inside both sidebands (Halpha, hemmed in by [N II] 6548 at -675 km/s and [N II] 6583 at +940, is the clearest case), and many more have one contaminated side. Because the same estimator is applied to prediction and reference, this does not displace the optimum – a perfect reconstruction still scores zero – but it makes the term degenerate: the model can match the measured flux by getting a line and its neighbour wrong in compensating directions. That is worst exactly for the blends, which are what deblending has to get right, so the answer is a better continuum rather than a shorter line list.

Pass all_positions (every catalogued line, (B,K)) to exclude sideband samples sitting within clean_half of a different line. Where fewer than min_clean samples survive, the plain median over the full sideband is used instead, so the estimate degrades to the old behaviour rather than to nothing.

Parameters:
Return type:

torch.Tensor

specsr.training.losses.line_presence_target(x_high, positions, *, in_range=None, valid=None, core_half=7, **mask_kw)[source]

Which catalogued lines the reference spectrum actually shows.

Returns (target, weight), both (B, K). target is 1 where the HR-derived line mask covers the line’s core window, and weight excludes lines that cannot be judged at all: those outside the observed grid, and those whose core the detector never measured. Excluding is important – scoring an unobservable line as a negative would train the presence head to switch off lines that are merely off the edge of the detector.

Deriving the label from the reference rather than from the line list is the same choice make_line_mask_from_smoothed() makes, for the same reason: it responds to lines present at the observed strength, not to every transition that happens to fall in range.

Parameters:
Return type:

tuple[torch.Tensor, torch.Tensor]

specsr.training.losses.sr1_deblend_loss(mean, log_var, x_high, x_high_err, valid=None, logvar_reg=3.38737824685736e-06, mask_smooth_k=121, mask_thresh_mad=7.5, mask_dilate=11, mask_min_width=7, lam_d1=0.11001162460004914, lam_d2=0.010153811105728492, gate_min_frac=0.015, gate_temp=0.05, score_w_recon=0.2, score_w_line=2.0, eps=1e-12)[source]

Heteroscedastic likelihood plus gated sharpness matching, for SR1.

Three parts:

  1. Likelihood. Gaussian NLL where the total variance is the sum of the predicted model variance and the reference spectrum’s own measurement variance. Keeping these separate matters: the model should not be rewarded for reporting confidence the data cannot support, nor penalised for failing to match noise.

  2. Gated sharpness. First and second derivatives of the prediction are matched to the reference, but only inside the line mask and normalised by the robust scale of the reference’s own derivatives. Matching derivatives is what forces narrow features to be reproduced at the right width — plain NLL is happy with a broadened line of the correct integrated flux.

  3. Gate. The sharpness term is scaled by a smooth function of how much of the spectrum the line mask covers. Spectra with essentially no detected lines contribute almost nothing, so featureless objects do not push the model to sharpen noise.

Returns the total loss and a dict of diagnostics for logging.

valid marks wavelengths the reference actually measured. Two things are required of the caller, and they are different:

  • Invalid samples must carry a neutral numerical value (zero, in the per-spectrum normalised units this loss works in) rather than nan or an arbitrary number. The line mask is derived from x_high by smoothing over a wide kernel, so a wild value at an unmeasured wavelength would be dragged into neighbouring measured samples and be mistaken for a line.

  • valid then removes those samples from the objective, so no gradient rewards matching the neutral value.

Both are needed. Filling alone is what the original pipeline did — it used the per-spectrum median and trained on it, teaching the model to emit flat continuum wherever the detector recorded nothing. Masking alone leaves the fill value free to contaminate any neighbourhood statistic.

Parameters:
Return type:

tuple[torch.Tensor, dict]

specsr.training.losses.sr2_loss(*, sr2_mean, sr2_logvar, x_high, x_high_err, line_mask, valid=None, presence, zhead=None, z_true=None, ztransform=None, use_sigma=True, lam_z=0.0, lam_hp_in=0.0, lam_hp_out=0.0, hp_k=51, lam_sparse=0.0, var_floor=1e-08, hp_in_noise_weighted=False, presence_logit=None, line_positions=None, line_in_range=None, lam_presence=0.0, lam_flux=0.0, presence_pos_weight=0.0, presence_pos_weight_max=50.0, flux_core_half=7, flux_sb_lo=11, flux_sb_hi=20, flux_eps=1.0, flux_beta=0.5, flux_weight_power=1.0, flux_clean_half=7, presence_mask_smooth_k=259, presence_mask_thresh_mad=7.5, presence_mask_dilate=11, presence_mask_min_width=5)[source]

Likelihood, high-pass matching, presence sparsity and a redshift term.

High-pass matching compares the small-scale content of prediction and reference after removing the smooth component. It is applied separately inside and outside the line mask (lam_hp_in / lam_hp_out) because the two regions want opposite behaviour: sharp structure should be reproduced on lines, and suppressed off them. A single weight cannot express that.

The two halves are weighted differently, and this is not an oversight. Outside the lines the term is inverse-variance weighted, which is the ordinary thing to do. Inside the lines it is not.

The reason is a comparison within the mask, not between lines and continuum: hp_in is a weighted mean over masked pixels alone, so a uniform rescaling of those weights cancels and the continuum’s weight is irrelevant to it. What survives is the spread inside the mask. Photon noise scales with flux, so the bright line cores are the noisiest pixels present and 1/sigma^2 drives their weight to near zero: on the validation split the brightest 10% of in-line pixels take 2.1% of the in-line weight, and the faintest 10% are weighted 34x more heavily than the brightest. The term that exists to fit lines was therefore decided by the faintest pixels in the mask and nearly blind to the flux-carrying cores, which starved the line branch and collapsed the presence gate.

Set hp_in_noise_weighted=True to restore the old behaviour; it is kept only so the ablation can be run.

Presence supervision (lam_presence) trains the gate as a classifier against the lines the reference spectrum actually shows, labelled by line_presence_target(). This replaces the older blanket sparsity penalty (lam_sparse, now off by default), which pushed every line logit down at a constant rate regardless of whether the line was real and drove the gate to a constant ~0.002 — the same value for a bright line as for one that is absent. Because the model forms amp * presence, that constant multiplied all line flux, and the line branch ended up supplying 0.17% of the flux the reference required. Sparsity is still what keeps spurious lines off, but it now comes from the negative half of a two-sided objective rather than from a uniform downward force.

The presence mask settings are grid-dependent and must be passed, which is why presence_mask_* are explicit parameters rather than left to make_line_mask_from_smoothed()’s defaults. Those defaults are sized for the retired 2,500-point linear grid, and this function silently used them until 2026-08-02 – so the mask defining both the presence labels and the flux term’s line list (w_flux = w_line * tgt) was built with the wrong kernel sizes on the log grid.

Measured on the 572-galaxy val split against integrated HR line SNR, the cost was carried almost entirely by min_width:

setting

recall@20

fpr@<3

precision

min_width=7

0.459

0.012

0.631

min_width=5

0.735

0.031

0.520

min_width=3

0.865

0.067

0.374

At 7 the mask missed 55% of the lines with HR SNR >= 20 – they were labelled absent for the presence BCE and dropped from the flux term. The cause is not the threshold: only 7.4% of those lines peak below 7.5 sigma per pixel (median 32.7). It is that keep_only_wide() demands min_width consecutive supra-threshold samples, and an R=1000 line on this R=4000 grid has sigma ~1.7 samples, so it clears 7.5 sigma over only ~5-6. smooth_k and dilate were measured to be near-neutral for this consumer over 121-430 and 11-29 respectively.

Do not “restore” the nominal log-grid rescaling recorded in by rescaling arithmetic (smooth_k=430, dilate=29, min_width=19): measured, it flags 1.1% of bright lines, because 19 samples is 1425 km/s and no real line is that wide.

Line flux matching (lam_flux) compares continuum-subtracted integrated flux in a velocity window around each real line, prediction against reference. Nothing else in this objective measures integrated line flux: the likelihood is inverse-variance weighted and so is quietest at the bright cores, and the high-pass term is a shape comparison that a correct profile at the wrong amplitude can still satisfy reasonably well.

Redshift coupling optionally back-propagates a redshift error through SR2, so the refinement is pushed towards spectra from which redshift is recoverable. The error is scaled by 1 / (1 + z), the standard convention, so high-redshift objects are not weighted disproportionately.

ztransform is a RedshiftTransform; passing it rather than loose statistics guarantees the decoding matches what the redshift head was trained with.

valid marks wavelengths the reference actually measured. Two things are required of the caller, and they are different:

  • Invalid samples must carry a neutral numerical value (zero, in the per-spectrum normalised units this loss works in) rather than nan or an arbitrary number. The line mask is derived from x_high by smoothing over a wide kernel, so a wild value at an unmeasured wavelength would be dragged into neighbouring measured samples and be mistaken for a line.

  • valid then removes those samples from the objective, so no gradient rewards matching the neutral value.

Both are needed. Filling alone is what the original pipeline did — it used the per-spectrum median and trained on it, teaching the model to emit flat continuum wherever the detector recorded nothing. Masking alone leaves the fill value free to contaminate any neighbourhood statistic.

Parameters:
Return type:

tuple[torch.Tensor, dict]