specsr.training.zhead_sources¶

Input representations for the redshift head.

The redshift experiment trains the same ZHead1D architecture on four different spectral representations, so that any difference in redshift accuracy is attributable to the information content of the input rather than to model capacity:

`name`	pipeline	channels
`lowres`	`flux_low` (prism) directly	1
`hires`	`flux_high` (grating) directly	1
`sr1`	LR -> SR1 -> (mean, log sigma)	2
`sr2`	LR -> SR1 -> SR2 -> (mean, log sigma)	2

Previously these lived in four near-duplicate copies of the training script (redshift_head, _hires, _lowres, _sr2) which had drifted apart by hundreds of lines, making it genuinely hard to confirm that the architecture and training procedure were held fixed across the comparison. Collapsing them to one script plus a swappable source makes that guarantee structural: there is only one training loop, one loss and one split.

hires is the upper bound of the comparison and lowres the lower bound; sr1/sr2 are the quantities of interest.

Functions

build_source(name, **kwargs)

Construct a source by name.

Classes

`RawFluxSource`(key, name[, n_channels])	Feed an observed spectrum straight to the head, with no model in between.
`SR1Source`(sr1[, use_sigma, name])	Run the frozen SR1 backbone and hand its output to the head.
`SR2Source`(sr1, zhead_bootstrap, sr2, ztransform)	Run the frozen SR1 -> ZHead -> SR2 chain and hand SR2's output to the head.
`ZHeadSource`(args, *kwargs)	Produces the `(B, C, L)` tensor fed to the redshift head.

class specsr.training.zhead_sources.ZHeadSource(*args, **kwargs)[source]¶

Bases: Protocol

Produces the (B, C, L) tensor fed to the redshift head.

name: str¶

n_channels: int¶

class specsr.training.zhead_sources.RawFluxSource(key, name, n_channels=1)[source]¶

Bases: object

Feed an observed spectrum straight to the head, with no model in between.

Used for the lowres and hires ends of the comparison.

Parameters:

key (str)
name (str)
n_channels (int)

key: str¶

name: str¶

n_channels: int = 1¶

class specsr.training.zhead_sources.SR1Source(sr1, use_sigma=True, name='sr1')[source]¶

Bases: object

Run the frozen SR1 backbone and hand its output to the head.

When use_sigma is set, the predictive log-sigma is supplied as a second channel so the head can down-weight wavelengths SR1 is unsure about. Note the conversion: SR1 emits log_var, and the head consumes log_sigma = 0.5 * log_var.

Parameters:

sr1 (torch.nn.Module)
use_sigma (bool)
name (str)

sr1: torch.nn.Module¶

use_sigma: bool = True¶

name: str = 'sr1'¶

property n_channels: int¶

class specsr.training.zhead_sources.SR2Source(sr1, zhead_bootstrap, sr2, ztransform, delta_cap=0.0, use_sigma=True, use_sr1_sigma=True, use_line_mask=True, use_zhat_channel=True, use_zsigma_channel=False, zsigma_line_mask=False, zsigma_mask_max_um=0.05, sigma_base_um=0.005, name='sr2')[source]¶

Bases: object

Run the frozen SR1 -> ZHead -> SR2 chain and hand SR2’s output to the head.

SR2 is conditioned on a redshift, so this needs a bootstrap redshift head (the one trained on the SR1 representation) to supply it. That head is frozen and is not the head being trained here — the model under evaluation only ever sees the final SR2 spectrum.

The SR2 input stack is built by specsr.models.sr2.build_sr2_input(), whose channel order and flags must match what the SR2 checkpoint was trained with. Take those flags from the checkpoint’s stored config rather than assuming defaults.

Parameters: