Reference / API reference / dorothy / dorothy.training / dorothy.training.batches

dorothy.training.batches

Batch construction for training and evaluation.

Every training mode (single-survey, dense multi-survey, sparse multi-survey, multi-labelset) is expressed as a BatchSource that yields Batch objects with the same structure, so the trainer runs a single train/validation loop:

• X: survey name -> spectra (batch, 3, n_wavelengths) = [flux, rho, mask]

• has_data: survey name -> (batch,) bool, which stars have that survey

• y: label source name -> labels (batch, 3, n_params) = [values, errors, mask]; stars without labels from a source have an all-zero mask.

Single-survey inputs use the survey key "default".

Functions

forward_label_sources(model, batch, label_keys)

Run the model on a batch, returning (batch, 2, n_params) per label source.

Classes

Batch(X, has_data, y)	One batch of model inputs and targets (see module docstring).
SparseBatchSource(data, indices, labels, device)	Batches built on the fly from memory-efficient sparse survey storage.
TensorBatchSource(X, y, has_data, device)	Batches from in-memory tensors (single- or multi-survey).

class dorothy.training.batches.Batch(X, has_data, y)

Bases: object

One batch of model inputs and targets (see module docstring).

class dorothy.training.batches.TensorBatchSource(X, y, has_data, device)

Bases: object

Batches from in-memory tensors (single- or multi-survey).

Parameters:

• X (dict[str, torch.Tensor]) – Survey name -> spectra tensor (n_samples, 3, n_wavelengths).

• y (dict[str, torch.Tensor]) – Label source name -> labels tensor (n_samples, 3, n_params).

• has_data (dict[str, torch.Tensor] | None) – Survey name -> (n_samples,) bool tensor. Defaults to all True.

• device (torch.device) – Device batches are moved to.

Raises:

ValueError – If a tensor of X, y or has_data has a different number of samples than y.

property label_keys: list[str]

Label sources this source provides targets for.

batches(batch_size, rng=None)

Yield batches, shuffled with rng if given, else in order.

Parameters:

• batch_size (int) – Number of samples per batch; the last batch may be smaller.

• rng (Generator | None) – Generator for the shuffle. None keeps the sample order.

Yields:

Batches of tensors on self.device.

class dorothy.training.batches.SparseBatchSource(data, indices, labels, device)

Bases: object

Batches built on the fly from memory-efficient sparse survey storage.

Spectra are stored only for stars that have them (see SparseMergedData); dense batches are assembled per step. Only the rows of stars that have a survey are gathered, moved to the device and given a mask; the other rows stay zero. On CUDA the gathered rows go through pinned memory with asynchronous copies, and nothing in batch assembly waits for the GPU.

Parameters:

• data (SparseMergedData) – Sparse survey data (spectra and index mappings).

• indices (np.ndarray) – Global star indices belonging to this source (e.g. the training split).

• labels (dict[str, np.ndarray]) – Label source name -> labels (n_total, 3, n_params), indexed by global star index (typically normalized copies of the data labels).

• device (torch.device) – Device batches are moved to.

property label_keys: list[str]

Label sources this source provides targets for.

build_batch(global_indices)

Assemble a dense batch for the given global star indices.

Each survey gives X (batch, 3, n_wavelengths) = [flux, rho, mask] with mask = 1 where rho > 0, and all-zero rows for stars without that survey.

Parameters:

global_indices (ndarray) – Global star indices of the batch, shape (batch,).

Returns:

The batch, with y from self.labels at these indices.

Return type:

Batch

batches(batch_size, rng=None)

Yield batches, shuffled with rng if given, else in order.

Parameters:

• batch_size (int) – Number of stars per batch; the last batch may be smaller.

• rng (Generator | None) – Generator for the shuffle. None keeps the order of indices.

Yields:

Batches built by build_batch().

dorothy.training.batches.forward_label_sources(model, batch, label_keys, label_errors=None)

Run the model on a batch, returning (batch, 2, n_params) per label source.

A multi-label MultiHeadMLP returns every head; a single-head model’s output is keyed by label_keys[0].

Parameters:

• model (MLP | MultiHeadMLP) – The model.

• batch (Batch) – The batch.

• label_keys (list[str]) – Label source names, in model output-head order.

• label_errors (dict[str, Tensor] | None) – Label source name -> label errors (batch, n_params) for that head’s log-scatter, keyed like the targets. None, or a head without an entry, uses the learned missing vector.

Returns:

Label source name -> output (batch, 2, n_params) = [mu, ln_s].

Return type:

dict[str, Tensor]