DOROTHY home

Reference / API reference / dorothy / dorothy.inference / dorothy.inference.predictor

dorothy.inference.predictor

Predict stellar parameters and uncertainties with trained DOROTHY models.

Predictor loads a checkpoint, predicts in chunks, denormalizes to physical units and applies the calibration factors to the predicted scatter.

Examples

>>> predictor = Predictor.load("path/to/checkpoint")
>>> result = predictor.predict(spectra, denormalize=True)
>>> print(result.predictions.shape, result.uncertainties.shape)

Functions

predict_from_fits(fits_path, checkpoint_path)

Predict stellar parameters for every spectrum in a FITS file.

Classes

PredictionResult(predictions, uncertainties, ...)

Container for prediction results.

Predictor(model[, normalizer, device, ...])

Predictor for stellar parameters from spectroscopic data.
class dorothy.inference.predictor.PredictionResult(predictions, uncertainties, raw_output, parameter_names, is_normalized=True, calibrated=False)[source]

Bases: object

Container for prediction results.

predictions

Predicted stellar parameters of shape (n_samples, n_params).

uncertainties

Predicted scatter, scaled by the calibration factors if calibrated, of shape (n_samples, n_params).

raw_output

Raw model output of shape (n_samples, 2, n_params) where raw_output[:, 0, :] contains predicted means and raw_output[:, 1, :] contains predicted log-scatter.

parameter_names

List of parameter names in order.

is_normalized

Whether predictions and uncertainties are in normalized units (False: physical units).

calibrated

Whether the uncertainties are scaled by calibration factors.

property n_samples: int

Number of predicted samples.

property n_parameters: int

Number of stellar parameters.

to_dict()[source]

Convert to columns by parameter name.

Returns:

Dict with <name> (prediction) and <name>_err (uncertainty) for each parameter, each of shape (n_samples,).

Return type:

dict[str, NDArray[float32]]

class dorothy.inference.predictor.Predictor(model, normalizer=None, device='cpu', parameter_names=None, scatter_floor=0.01, survey=None, label_source=None, calibration=None)[source]

Bases: object

Predictor for stellar parameters from spectroscopic data.

This class provides inference capabilities for trained DOROTHY models (single-survey MLP or multi-survey MultiHeadMLP), including loading checkpoints, making predictions, and denormalizing results to physical units.

Inputs use the training format: arrays of shape (n_samples, 3, n_wavelengths) with channels [flux, rho, mask], where rho is the precision encoding of sigma (encode_precision) unless the checkpoint’s data.precision_encoding is false. For a MultiHeadMLP, pass either one survey’s array (set survey) or a dict mapping survey names to arrays (optionally with has_data masks).

model

The trained neural network model.

normalizer

Label normalizer for denormalization (optional).

device

Device to run inference on.

parameter_names

Names of stellar parameters being predicted.

scatter_floor

Scatter floor used to convert log-scatter to scatter.

survey

Survey of single-array inputs (MultiHeadMLP only).

label_source

Output head to predict from (MultiHeadMLP only; None uses the first head).

calibration

Calibration factor per parameter name of the head predicted from, which scales the predicted scatter (see calibrate_scatter()), or None for the raw scatter.

config

Experiment config of the checkpoint, if loaded from one.

Examples

>>> predictor = Predictor.load("checkpoint/")
>>> result = predictor.predict(X, denormalize=True)
>>> print(result.predictions.shape)  # (n_samples, n_params)
classmethod load(checkpoint_path, model_file='best_model.pth', normalizer_file='normalizer.pkl', device='auto', survey=None, label_source=None, calibrate=True)[source]

Load a predictor from a checkpoint directory.

See load_checkpoint() for how the model is rebuilt.

Parameters:

  • checkpoint_path (str | Path) – Path to checkpoint directory.

  • model_file (str) – Name of model weights file.

  • normalizer_file (str | None) – Name of normalizer file (None to skip).

  • device (str) – Device for inference (“auto”, “cpu”, “cuda”, “mps”).

  • survey (str | None) – For a MultiHeadMLP, the survey of single-array inputs.

  • label_source (str | None) – For a MultiHeadMLP, which output head to use.

  • calibrate (bool) – Scale the predicted scatter by the checkpoint’s calibration factors for that head, if it has them.

Returns:

Configured Predictor instance.

Raises:

  • FileNotFoundError – If the weights file does not exist.

  • ValueError – If a checkpoint without config.yaml holds a model with label-aware uncertainty, or survey or label_source is invalid for the model.

Return type:

Predictor

classmethod from_config(config, checkpoint_path=None)[source]

Create a predictor from experiment configuration.

Parameters:

  • config (ExperimentConfig) – Experiment configuration.

  • checkpoint_path (str | Path | None) – Path to checkpoint (uses config default if None).

Returns:

Configured Predictor instance, on config.device.

Raises:

FileNotFoundError – If the weights file does not exist.

Return type:

Predictor

property is_multi_head: bool

Whether the model is a MultiHeadMLP.

predict(X, batch_size=1024, denormalize=False, has_data=None)[source]

Make predictions on input spectra.

Parameters:

  • X (NDArray[float32] | Tensor | dict[str, NDArray[float32]]) – Input spectra of shape (n_samples, 3, n_wavelengths), or for a MultiHeadMLP a dict mapping survey names to such arrays.

  • batch_size (int) – Batch size for inference.

  • denormalize (bool) – Convert predictions and uncertainties to physical units. Without a normalizer they stay normalized.

  • has_data (dict[str, NDArray[bool]] | None) – For dict inputs, survey name -> (n_samples,) bool array of which stars have that survey (default: all).

Returns:

PredictionResult containing predictions and uncertainties – the predicted scatter, times the calibration factors if the predictor has them.

Raises:

ValueError – If the calibration factors are not one per parameter name.

Return type:

PredictionResult

predict_chunked(X_iterator, denormalize=False, batch_size=1024)[source]

Make predictions on data provided by an iterator.

This is memory-efficient for large datasets that don’t fit in memory.

Parameters:

  • X_iterator (Iterator[NDArray[float32]]) – Iterator yielding batches of input spectra.

  • denormalize (bool) – Whether to denormalize predictions.

  • batch_size (int) – Batch size for inference within each chunk.

Yields:

PredictionResult for each batch.

predict_all_chunked(X_iterator, denormalize=False, batch_size=1024)[source]

Make predictions on all data from iterator and concatenate.

Parameters:

  • X_iterator (Iterator[NDArray[float32]]) – Iterator yielding batches of input spectra.

  • denormalize (bool) – Whether to denormalize predictions.

  • batch_size (int) – Batch size for inference within each chunk.

Returns:

Combined PredictionResult for all batches.

Return type:

PredictionResult

get_embeddings(X, layer_index=-2, batch_size=1024)[source]

Extract intermediate embeddings from the model.

Useful for k-NN anomaly detection and dimensionality reduction. For an MLP, layer_index selects the hidden layer. For a MultiHeadMLP, the shared trunk features of self.survey (comparable across surveys) are returned and layer_index is ignored.

Parameters:

  • X (NDArray[float32] | Tensor) – Input spectra of shape (n_samples, 3, n_wavelengths).

  • layer_index (int) – Which MLP layer to extract embeddings from.

  • batch_size (int) – Batch size for processing.

Returns:

Embeddings array of shape (n_samples, embedding_dim).

Raises:

ValueError – If the model is a MultiHeadMLP with several surveys and survey is not set.

Return type:

NDArray[float32]

dorothy.inference.predictor.predict_from_fits(fits_path, checkpoint_path, output_path=None, chunk_size=1000, device='auto', batch_size=1024, denormalize=True, survey=None, label_source=None, calibrate=True)[source]

Predict stellar parameters for every spectrum in a FITS file.

Spectra are preprocessed as in training: per-spectrum flux normalization when the checkpoint’s data.normalize_spectra is set, the 3-channel input format, and the precision encoding of sigma when data.precision_encoding is set. Without a config.yaml, spectra are normalized and not encoded.

Parameters:

  • fits_path (str | Path) – Path to input FITS file (SPECTRA / SPEC_IVAR extensions).

  • checkpoint_path (str | Path) – Path to model checkpoint directory.

  • output_path (str | Path | None) – Path to save predictions (CSV): the star ID, each parameter and its uncertainty, and a calibrated column that says whether the uncertainties are calibrated. If None, returns results.

  • chunk_size (int) – Number of spectra to process at a time.

  • device (str) – Device for inference.

  • batch_size (int) – Batch size for inference.

  • denormalize (bool) – Output physical units (requires a normalizer).

  • survey (str | None) – Survey of the spectra, for multi-survey models.

  • label_source (str | None) – Output head to use, for a MultiHeadMLP.

  • calibrate (bool) – Scale the uncertainties by the checkpoint’s calibration factors, if it has them.

Returns:

PredictionResult if output_path is None, otherwise None.

Raises:

ValueError – If the spectra do not match the model’s wavelength grid, or the model has several surveys and survey is not set.

Return type:

PredictionResult | None