specsr.training.ztransform

Redshift normalisation and bounded decoding.

The redshift head does not regress z directly. Training normalises z to zero mean and unit variance using statistics from the training split only, and the head’s raw output is squashed through a sigmoid onto the observed redshift range before the loss is applied.

The sigmoid bound matters: an unbounded head can emit redshifts far outside the range the data covers, and the heteroscedastic likelihood will happily trade a wild prediction against a large predicted variance. Bounding the mean keeps the uncertainty interpretable. It also keeps gradients alive at the extremes, which a hard clamp would not.

Collecting this in one object matters because the same transform must be applied identically at training, inference and evaluation time. When it lived inline in four copies of the training script, a divergence would have shown up as a subtle redshift bias rather than as an error.

Classes

RedshiftTransform(mean, std, z_min_n, z_max_n)

Affine normalisation plus a sigmoid bound on the decoded mean.
class specsr.training.ztransform.RedshiftTransform(mean, std, z_min_n, z_max_n)[source]

Bases: object

Affine normalisation plus a sigmoid bound on the decoded mean.

Variables:

  • std (mean,) – Normalisation statistics, computed on the training split only so no information about held-out redshifts leaks through the transform.

  • z_max_n (z_min_n,) – Normalised bounds of the observed redshift range, used to squash the head’s raw output.

Parameters:
mean: float
std: float
z_min_n: float
z_max_n: float
classmethod from_training_redshifts(z_train, pad=0.0)[source]

Fit the transform to the training-split redshifts.

pad optionally widens the bounds by a fraction of the range, giving the sigmoid a little headroom at the extremes.

Parameters:

pad (float)

Return type:

RedshiftTransform

normalize(z)[source]

Physical redshift -> normalised units.

Parameters:

z (torch.Tensor)

Return type:

torch.Tensor

denormalize(z_n)[source]

Normalised units -> physical redshift.

Parameters:

z_n (torch.Tensor)

Return type:

torch.Tensor

decode_mean(mu_raw, bounded=True)[source]

Squash a raw head output onto the observed range, in normalised units.

bounded=False passes the value through untouched, for heads whose estimate is already in normalised redshift units and already inside the range – the classification head reads its estimate off a grid, so squashing it a second time would compress it towards the range centre. Pass zhead.bounded_mean rather than deciding per call site.

Parameters:
Return type:

torch.Tensor

decode_sigma(log_var, var_floor=1e-06)[source]

Predicted log-variance (normalised) -> physical 1-sigma uncertainty.

Parameters:
Return type:

torch.Tensor

predict(mu_raw, log_var, var_floor=1e-06, bounded=True)[source]

Raw head outputs -> physical (z, sigma_z).

log_var is clamped before exponentiation; without it an early, overconfident-in-the-wrong-direction update can send exp(log_var) to infinity and take the loss with it.

Parameters:
Return type:

tuple[torch.Tensor, torch.Tensor]

as_dict()[source]
Return type:

dict

classmethod from_dict(d)[source]
Parameters:

d (dict)

Return type:

RedshiftTransform