specsr.inference.pipeline

The user-facing inference pipeline: LR spectrum in, SR spectrum and redshift out.

This is the object the README promises. It wraps the three stages so a caller never has to know the channel order of the SR2 input stack, that SR1 emits log_var while the head consumes log_sigma, or that the redshift decoding has to use the transform the head was trained with. Getting any of those subtly wrong produces a plausible-looking spectrum rather than an error.

Everything returned is in the physical units of the input spectrum, not the per-spectrum normalised units the models work in. The normalisation is applied on the way in and undone on the way out, using the input’s own moments.

Functions

resolve_zhead_in(directory)

Find the redshift head in a checkpoint directory, under either name.

run_infer(args)

CLI entry point for specsr infer.

Classes

SpecSRPipeline(sr1, zhead, sr2, cfg, ...)

Frozen SR1 -> redshift head -> SR2, ready for inference.

SpecSRResult(sr1, sr1_sigma, sr2, sr2_sigma, ...)

Output of one pipeline call, in the input's physical flux units.
class specsr.inference.pipeline.SpecSRPipeline(sr1, zhead, sr2, cfg, ztransform, wave, device)[source]

Bases: object

Frozen SR1 -> redshift head -> SR2, ready for inference.

Construct with from_checkpoints(). Calling the pipeline on a batch of low-resolution spectra returns a SpecSRResult.

property wavelength: ndarray

The model’s output grid in microns, shape (L_hi,).

Alias of wave, matching SpecSRResult.wavelength so the same name means the same thing on both objects.

classmethod from_checkpoints(directory=None, *, sr1_ckpt=None, sr1_config=None, zhead_ckpt=None, sr2_ckpt=None, wavelength=None, dataset=None, device=None)[source]

Load a chain, either from a directory of checkpoints or explicit paths.

A directory is expected to hold best_superres_model.pth, config_logR.yaml, a redshift head and optionally best_sr2.pth – the layout the archived run directories use.

The head is accepted under either name: best_zhead.pth (the Hub and bundle layout) or best_zhead_<source>.pth (what specsr train zhead actually writes). Only the first was recognised until 2026-08-14, so a directory produced directly by the trainer could not be loaded here at all – which is the normal case when evaluating a run you have just finished.

The HR wavelength grid must come from somewhere: pass wavelength directly, or dataset to read it from a built .npz.

Parameters:

  • directory (str | Path | None)

  • sr1_ckpt (str | Path | None)

  • sr1_config (str | Path | None)

  • zhead_ckpt (str | Path | None)

  • sr2_ckpt (str | Path | None)

  • wavelength (np.ndarray | None)

  • dataset (str | Path | None)

  • device (str | torch.device | None)

Return type:

SpecSRPipeline

classmethod from_pretrained(repo_id=None, revision=None, *, dataset=None, wavelength=None, device=None)[source]

Load the published chain, fetching weights from the Hugging Face Hub.

Set SPECSR_CHECKPOINT_DIR to load from a local directory instead – no network and no Hub account. See specsr.checkpoints.

The default revision deliberately excludes v1-submission: those weights were trained on a leaky split and reproduce numbers that are not real.

Parameters:

  • repo_id (str | None)

  • revision (str | None)

  • dataset (str | Path | None)

  • wavelength (np.ndarray | None)

  • device (str | torch.device | None)

Return type:

SpecSRPipeline

class specsr.inference.pipeline.SpecSRResult(sr1, sr1_sigma, sr2, sr2_sigma, z, z_sigma, wavelength)[source]

Bases: object

Output of one pipeline call, in the input’s physical flux units.

Parameters:
sr1: ndarray

SR1 reconstruction, shape (B, L_hi).

sr1_sigma: ndarray

SR1 predictive 1-sigma.

sr2: ndarray | None

SR2 refinement, or None if no SR2 checkpoint was loaded.

sr2_sigma: ndarray | None
z: ndarray

Predicted redshift, shape (B,).

z_sigma: ndarray
wavelength: ndarray

HR wavelength grid in microns, shape (L_hi,).

specsr.inference.pipeline.resolve_zhead_in(directory)[source]

Find the redshift head in a checkpoint directory, under either name.

Two naming conventions coexist, and both are load-bearing. best_zhead.pth is what the Hub layout and the assembled bundles use. best_zhead_<source>.pth is what specsr train zhead writes, because the four comparison arms of the redshift experiment are distinguished by that suffix and three of them are published under it (zhead/best_zhead_lowres.pth and friends in specsr.checkpoints).

Renaming either would break the other, so lookup accepts both – specsr.checkpoints already did this via _LOCAL_ALIASES and this is the same policy for directory-based loading.

Raises rather than guessing when a directory holds several arms: choosing one by sort order would silently pair SR2 with the wrong redshift head, and every line position it emits depends on that choice.

Parameters:

directory (Path)

Return type:

Path | None

specsr.inference.pipeline.run_infer(args)[source]

CLI entry point for specsr infer.

Return type:

int