Reference / API reference / dorothy / dorothy.training / dorothy.training.trainer

dorothy.training.trainer¶

Training infrastructure for DOROTHY models.

This module implements the training loop for stellar parameter prediction, including: - Model and optimizer setup from configuration - Batch-wise training with gradient clipping - Learning rate scheduling (CyclicLR, ReduceOnPlateau, etc.) - A β-NLL warm-up of the loss - An optional exponential moving average of the weights - Validation evaluation and best model tracking - Checkpointing and training history - Early stopping support

All training modes (single-survey, dense multi-survey, sparse multi-survey and multi-labelset) share one pipeline: inputs are wrapped in a batches source, labels are normalized by one routine, and a single epoch loop handles training, validation, scheduling, best-weight tracking, early stopping and history.

Classes

`Trainer`(config[, parameter_names])	Trainer for DOROTHY stellar parameter models.
`TrainingHistory`([train_losses, val_losses, ...])	Container for training metrics and history.

class dorothy.training.trainer.TrainingHistory(train_losses=<factory>, val_losses=<factory>, learning_rates=<factory>, nll_warmup_epochs=0.0, best_epoch=0, best_val_loss=inf, total_time=0.0, parameter_names=<factory>, val_loss_breakdown=<factory>, val_metrics=<factory>, weight_norms=<factory>, grad_norms=<factory>, weight_updates=<factory>, survey_names=<factory>, per_survey_val_losses=<factory>, per_survey_val_metrics=<factory>, label_source_names=<factory>, per_labelset_val_losses=<factory>, per_labelset_val_metrics=<factory>, per_labelset_best_epoch=<factory>, per_labelset_best_val_loss=<factory>)[source]¶

Bases: object

Container for training metrics and history.

train_losses¶: Per-epoch training loss values.

val_losses¶: Per-epoch validation loss values. With several label sources this combines all heads (see training.label_source_weighting) and selects the best epoch.

learning_rates¶: Learning rate after each optimizer step.

nll_warmup_epochs¶: Length of the β-NLL warm-up in epochs (warm-up steps divided by steps per epoch, so it can be fractional); 0 if there was no warm-up. During the warm-up the training loss is β-weighted and not comparable with the validation loss.

best_epoch¶: Epoch with the best validation loss (0-based).

best_val_loss¶: Best validation loss achieved.

total_time¶: Total training time in seconds.

parameter_names¶: Names of stellar parameters being tracked.

val_loss_breakdown¶: Per-epoch loss breakdown (primary label source) with keys: - ‘mean_component’: (n_epochs, n_params) (mu - y)^2 / s^2 per param - ‘scatter_component’: (n_epochs, n_params) log(s^2) per param

val_metrics¶: Per-epoch Evaluator metrics (primary label source) as dict of lists of arrays. Keys are metric names (rmse, bias, sd, mae, median_offset, robust_scatter, z_median, z_robust_scatter, pred_unc_p16, pred_unc_p50, pred_unc_p84). Each value is a list of arrays with shape (n_params,) per epoch.

weight_norms¶: Per-layer L2 weight norms per epoch for grokking detection. Keys are layer names, values are lists of norms per epoch.

grad_norms¶: Total gradient norm per epoch (average across batches).

weight_updates¶: Per-layer weight update magnitudes per epoch. Keys are layer names, values are lists of ||W_new - W_old|| per epoch.

survey_names / per_survey_val_losses / per_survey_val_metrics: Per-survey validation loss and metrics (primary label source).

label_source_names / per_labelset_val_losses / per_labelset_val_metrics: Per-head validation loss and metrics for multi-labelset training.

per_labelset_best_epoch / per_labelset_best_val_loss: Epoch at which each head’s own validation loss was lowest, and that loss.

save(path)[source]¶

Save the training history to a pickle file.

The pickle holds a dict of the fields, except learning_rates; train_losses and val_losses are stored as history_train and history_val.

Parameters:: path (str | Path) – File to write.

save_learning_rates(path)[source]¶

Save the learning rate of every optimizer step to a pickle file.

Parameters:: path (str | Path) – File to write; it holds the list learning_rates.

get_loss_breakdown_array()[source]¶

Get loss breakdown as numpy arrays.

Returns:: Dictionary with keys ‘mean_component’ and ‘scatter_component’, each with shape (n_epochs, n_params).
Return type:: dict[str, ndarray]

get_metrics_array()[source]¶

Get Evaluator metrics as numpy arrays.

Returns:: Dictionary with metric names as keys (rmse, bias, sd, mae, etc.), each with shape (n_epochs, n_params).
Return type:: dict[str, ndarray]

get_grokking_metrics()[source]¶

Get grokking detection metrics as numpy arrays.

Returns:: Dictionary containing – - ‘weight_norms’: dict mapping layer names to (n_epochs,) arrays - ‘grad_norms’: (n_epochs,) array of average gradient norms - ‘weight_updates’: dict mapping layer names to (n_epochs,) arrays
Return type:: dict[str, ndarray | dict[str, ndarray]]

class dorothy.training.trainer.Trainer(config, parameter_names=None)[source]¶

Bases: object

Trainer for DOROTHY stellar parameter models.

Handles the complete training pipeline including model creation, optimization, scheduling, validation, checkpointing, and label normalization.

Training Modes:

The trainer supports four training modes for different data configurations:

Single-Survey (fit): Standard MLP with one spectral survey. Use when training on a single survey (e.g., BOSS only).
Multi-Survey (fit_multi_survey): MultiHeadMLP with multiple surveys sharing a common label source. Use when combining spectra from different surveys (e.g., BOSS + DESI) with labels from one source.
Multi-Survey Sparse (fit_multi_survey_sparse): Same as multi-survey but uses memory-efficient sparse storage, with one or several label sources. This is the path used by dorothy train.
Multi-Labelset (fit_multi_labelset): MultiHeadMLP with multiple surveys AND multiple label sources from dense arrays.

All modes run the same pipeline, so they normalize labels, weight label sources, validate and select the best epoch identically.

Training Flow:

Initialize Trainer with ExperimentConfig
Trainer creates model, loss function, optimizer
Call fit*() with training and validation data
fit*() normalizes labels, creates the scheduler and runs the loop
Each epoch: train epoch -> validate -> update history
Best model weights saved based on validation loss
After training: calibration factors from the best epoch’s validation z-scores
save_checkpoint() persists the model, normalizer and factors

config¶: Experiment configuration.

model¶: The neural network model (MLP or MultiHeadMLP).

optimizer¶: The optimizer (AdamW by default).

scheduler¶: Learning rate scheduler (OneCycleLR by default).

loss_fn¶: Loss function (HeteroscedasticLoss, or MaskedMSELoss for loss: mse).

device¶: Compute device (cuda/mps/cpu).

autocast_dtype¶: Dtype of the training forward pass with mixed precision (training.amp on CUDA), or None for float32.

scaler¶: Gradient scaler, a pass-through unless the forward pass runs in float16.

history¶: TrainingHistory with loss curves and metrics.

normalizer¶: LabelNormalizer for converting to/from normalized space.

parameter_names¶: Names of stellar parameters being predicted.

calibration¶: Calibration factors fit after training, label source -> parameter name -> factor, or None when training.calibrate_uncertainties is false, the loss has no scatter, or training has not run.

Examples

>>> from dorothy.config import ExperimentConfig, DataConfig
>>> config = ExperimentConfig(
...     name="test",
...     data=DataConfig(catalogue_path=Path("/data/catalogue.h5")),
... )
>>> trainer = Trainer(config)
>>> # X_train shape: (n_samples, 3, n_wavelengths) - [flux, rho, mask]
>>> # y_train shape: (n_samples, 3, n_params) - [values, errors, mask]
>>> trainer.fit(X_train, y_train, X_val, y_val)
>>> trainer.save_checkpoint()  # Saves model + normalizer

property n_parameters: int¶: Get number of stellar parameters from config.

property label_keys: list[str]¶: Label source names, in model output-head order.

fit(X_train, y_train, X_val, y_val, normalize_labels=True)[source]¶

Train the model on a single survey.

Parameters:

X_train (ndarray | Tensor) – Training features of shape (n_samples, 3, n_wavelengths) with 3 channels: [flux, rho, mask].
y_train (ndarray | Tensor) – Training labels of shape (n_samples, 3, n_params) where: - Channel 0: label values - Channel 1: label errors (uncertainties) - Channel 2: mask, valid where > 0
X_val (ndarray | Tensor) – Validation features (same format as X_train).
y_val (ndarray | Tensor) – Validation labels (same format as y_train).
normalize_labels (bool) – Whether to fit and apply label normalization. If True, creates a LabelNormalizer, fits on training labels, and transforms both train and validation labels to normalized space. Normalization is applied to channels 0-1 (values and errors), channel 2 (mask) is preserved unchanged.

Returns:

TrainingHistory with loss curves and metrics.

Raises:

ValueError – If a label array does not have shape (n_samples, 3, n_params), the model has several surveys, or no label is valid for normalization.
RuntimeError – If no epoch has a finite validation loss.

Return type:

TrainingHistory

fit_multi_survey(X_train, y_train, X_val, y_val, has_data_train, has_data_val, normalize_labels=True)[source]¶

Train the model with multi-survey data.

This method handles training with multiple spectral surveys where each star may have data from one or more surveys. Uses MultiHeadMLP with survey-specific encoders.

Parameters:

X_train (dict[str, ndarray | Tensor]) – Dict mapping survey names to training spectral data arrays. Each array has shape (n_samples, 3, n_wavelengths) with 3 channels: [flux, rho, mask]. Missing data is indicated by mask=0.
y_train (ndarray | Tensor) – Training labels of shape (n_samples, 3, n_params) where: - Channel 0: label values - Channel 1: label errors (uncertainties) - Channel 2: mask, valid where > 0
X_val (dict[str, ndarray | Tensor]) – Dict mapping survey names to validation spectral data.
y_val (ndarray | Tensor) – Validation labels (same format as y_train).
has_data_train (dict[str, ndarray | Tensor]) – Dict mapping survey names to boolean arrays indicating which samples have data from that survey.
has_data_val (dict[str, ndarray | Tensor]) – Same for validation set.
normalize_labels (bool) – Whether to fit and apply label normalization.

Returns:

TrainingHistory with loss curves and metrics.

Raises:

TypeError – If model is not a MultiHeadMLP.
ValueError – If a label array does not have shape (n_samples, 3, n_params), or no label is valid for normalization.
RuntimeError – If no epoch has a finite validation loss.

Return type:

TrainingHistory

fit_multi_survey_sparse(data, train_indices, val_indices, normalize_labels=True)[source]¶

Train the model with memory-efficient sparse multi-survey data.

This method is designed for large multi-survey datasets where storing dense arrays for all surveys would exceed available memory. It builds dense batches on the fly from sparse storage, which holds only the stars that have each survey.

For a model with several output heads, the labels of each source come from data.labels_dict; otherwise data.labels is used.

Parameters:

data (SparseMergedData) – SparseMergedData from CatalogueLoader.load_merged_sparse(). Contains sparse per-survey spectra and dense labels.
train_indices (ndarray) – Global indices for training set.
val_indices (ndarray) – Global indices for validation set.
normalize_labels (bool) – Whether to fit and apply label normalization.

Returns:

TrainingHistory with loss curves and metrics.

Raises:

TypeError – If model is not a MultiHeadMLP.
ValueError – If the model’s label sources are not in the data, or no training label is valid for normalization.
RuntimeError – If no epoch has a finite validation loss.

Return type:

TrainingHistory

Examples

>>> loader = CatalogueLoader("super_catalogue.h5")
>>> data = loader.load_merged_sparse(["boss", "desi", "lamost_lrs"])
>>> train_idx, val_idx = ...  # global star indices of the two splits
>>> history = trainer.fit_multi_survey_sparse(data, train_idx, val_idx)

fit_multi_labelset(X_train, y_train, X_val, y_val, has_data_train, has_data_val, has_labels_train=None, has_labels_val=None, normalize_labels=True)[source]¶

Train the model with multi-survey inputs and multi-labelset outputs.

This method handles training where: - Different surveys provide spectral data (multi-survey inputs) - Different label sources provide targets (multi-labelset outputs) - Stars may have data from any subset of surveys and label sources

The approach uses union-style masking: if a star has labels from APOGEE but not GALAH, the loss for GALAH predictions is masked for that star.

Parameters:

X_train (dict[str, ndarray | Tensor]) – Dict mapping survey names to training spectral data arrays. Each array has shape (n_samples, 3, n_wavelengths) with 3 channels: [flux, rho, mask]. Missing data is indicated by mask=0.
y_train (dict[str, ndarray | Tensor]) – Dict mapping label source names to training label arrays. Each array has shape (n_samples, 3, n_params) where: - Channel 0: label values - Channel 1: label errors (uncertainties) - Channel 2: mask, valid where > 0
X_val (dict[str, ndarray | Tensor]) – Dict mapping survey names to validation spectral data.
y_val (dict[str, ndarray | Tensor]) – Dict mapping label source names to validation label arrays.
has_data_train (dict[str, ndarray | Tensor]) – Dict mapping survey names to boolean arrays indicating which samples have data from that survey.
has_data_val (dict[str, ndarray | Tensor]) – Same for validation set.
has_labels_train (dict[str, ndarray | Tensor] | None) – Dict mapping label source names to boolean arrays indicating which samples have labels from that source. If None, derived from y_train mask channels (any valid label).
has_labels_val (dict[str, ndarray | Tensor] | None) – Same for validation set. If None, derived from y_val mask channels.
normalize_labels (bool) – Whether to fit and apply label normalization. The normalizer is fit on the union of all label sources (the first source’s value is used where several have one).

Returns:

TrainingHistory with loss curves and metrics.

Raises:

TypeError – If model is not a MultiHeadMLP with multi-label support.
ValueError – If the label sources don’t match the model’s heads, a label array does not have shape (n_samples, 3, n_params), or no label is valid for normalization.
RuntimeError – If no epoch has a finite validation loss.

Return type:

TrainingHistory

save_checkpoint(path=None)[source]¶

Save the final model and training artifacts.

Saves:

best_model.pth: Best model weights (also loaded in self.model)
final_model.pth: Weights after the last training epoch
calibration.json: The calibration factors of best_model.pth, if fit; otherwise a calibration.json already in the directory is removed
normalizer.pkl: Label normalizer (if used)
history_train_val.pkl: Training/validation losses
learning_rates.pkl: Learning rate history
config.yaml: Experiment configuration

With training.ema_decay, both weight files hold averaged weights.

Parameters:: path (str | Path | None) – Directory to save to. If None, uses config output path.
Returns:: Path to the saved checkpoint directory.
Return type:: Path

predict(X, batch_size=2048)[source]¶

Make predictions with the trained model.

The model reads the learned missing vector in place of label errors.

Parameters:

X (ndarray | Tensor) – Input spectra of shape (n_samples, 3, n_wavelengths) = [flux, rho, mask], for a model with one survey.
batch_size (int) – Number of samples per forward pass.

Returns:

Array of shape (n_samples, 2, n_params) = [mu, ln_s], in normalized label units.

Raises:

ValueError – If the model has several surveys.

Return type:

ndarray

get_predictions_and_uncertainties(X)[source]¶

Get predictions and uncertainty estimates.

The scatter is meaningful only for a model trained with HeteroscedasticLoss.

Parameters:: X (ndarray | Tensor) – Input spectra, as for predict().
Returns:: Tuple of (predictions, uncertainties) as numpy arrays, both in normalized label units. predictions has shape (n_samples, n_params). uncertainties has shape (n_samples, n_params): the predicted scatter, without the calibration factors.
Raises:: ValueError – If the model has several surveys.
Return type:: tuple[ndarray, ndarray]

predict_multi_labelset(X, has_data)[source]¶

Make predictions from all label source output heads.

This method is used for inference after training with fit_multi_labelset(). It returns predictions from all label source heads.

Parameters:

X (dict[str, ndarray | Tensor]) – Dict mapping survey names to spectral data arrays. Each array has shape (n_samples, 3, n_wavelengths) = [flux, rho, mask].
has_data (dict[str, ndarray | Tensor]) – Dict mapping survey names to boolean arrays indicating which samples have data from that survey.

Returns:

Dict mapping label source names to prediction arrays of shape (n_samples, 2, n_params), in normalized label units: predicted means in [:, 0, :] and predicted log-scatter in [:, 1, :].

Raises:

TypeError – If model is not a MultiHeadMLP with multi-label support.

Return type:

dict[str, ndarray]