specsr.data.datasets

Torch datasets over the paired spectra products.

Each row is one (low-resolution prism, medium-resolution grating) pair on a shared wavelength grid, plus the reference uncertainty, the catalogue redshift and provenance.

Normalisation is per spectrum: each spectrum is standardised by its own mean and standard deviation. That is deliberate — it carries no information between rows, so unlike a dataset-wide normalisation it cannot leak anything about the held-out split. The reference uncertainty is divided by the same scale so it stays in the normalised units the loss works in, and the per-row mean and scale are returned so predictions can be pushed back into physical units for evaluation and plotting.

Functions

normalize_spectrum(x[, eps])

Standardise one spectrum by its own mean and standard deviation.

Classes

FixedGridSpectraDataset(*args, **kwargs)

Paired spectra on a fixed grid, returned as a tuple.

PairedSpectra(*args, **kwargs)

Paired low/high-resolution spectra from a built .npz product.
class specsr.data.datasets.FixedGridSpectraDataset(*args, **kwargs)[source]

Bases: Dataset

Paired spectra on a fixed grid, returned as a tuple.

(low, high, high_err, z, high_mean, high_std) — the last two let a caller de-normalise back into physical HR units, which is what line fluxes and residuals mean anything in.

This is the loader the training scripts use, and it lives here rather than in train/ so the normalisation exists once. It previously had a second copy in train/sr1_best/train_sr1.py; the HR-error fix below then had to be applied twice, which is exactly the kind of duplication that lets one copy quietly keep a bug.

PairedSpectra is the dict-returning equivalent and the better API for new code; this exists because the trained checkpoints were produced against the tuple form.

Parameters:

  • normalize_flux (bool)

  • target_key_raw (str)

class specsr.data.datasets.PairedSpectra(*args, **kwargs)[source]

Bases: Dataset

Paired low/high-resolution spectra from a built .npz product.

Parameters:

  • npz_path (str | Path) – Path to a product written by the preprocessing pipeline.

  • normalize_flux (bool) – Standardise each spectrum by its own statistics (default). When off, fluxes are returned as stored but the per-row statistics are still computed and returned, so de-normalisation helpers keep working.

  • target_key (str) – Which reference to train against — "flux_high" for the stitched grating spectrum, or "flux_high_smoothed" for its smoothed variant. The matching *_err array is used as the reference uncertainty.

  • dtype (torch.dtype) – Storage dtype. float32 halves memory against the float64 the products are written in, which matters: at the log wavelength grid (~6.7k samples) a 52k-row product is several GB per array.

Returns:

  • Each item is a dict so that call sites index by name rather than by

  • position. The previous tuple-based interface made it easy to silently

  • transpose two fields at a call site; a dict makes that a KeyError.

  • flux_low, flux_high, flux_high_err are (L,) tensors;

  • z is a scalar tensor; flux_high_mean/flux_high_std are the

  • per-row de-normalisation statistics; index is the row index.

REQUIRED = ('flux_low', 'z')

Arrays required to be present in a product.

property n_samples: int

Length of the wavelength axis.

denormalize(x, idx)[source]

Map normalised flux back to the physical units of the reference.

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

specsr.data.datasets.normalize_spectrum(x, eps=1e-25)[source]

Standardise one spectrum by its own mean and standard deviation.

Returns (normalised, mean, std). NaN-aware, and the scale is floored so a flat or empty spectrum cannot produce a division by zero.

Parameters:
Return type:

tuple[ndarray, float, float]