DOROTHY home

Reference / API reference / dorothy / dorothy.data / dorothy.data.training_data

dorothy.data.training_data

Load, filter and split experiment data from a config.

dorothy train and dorothy evaluate both go through load_training_data() and DataSplit, so the stars used for training and the held-out test set are defined in exactly one place.

Three layouts are produced, chosen by the config:

  • single_survey: one survey and one label source, dense arrays X (N, 3, W) and y (N, 3, P), for a plain MLP (or a single-survey MultiHeadMLP).

  • dense_multi_label: several label sources as dense dicts (data.use_dense_loading: true; high memory).

  • sparse_multi_survey: every other configuration, in memory-efficient SparseMergedData (the default).

Rows are the loaded arrays’ rows (global catalogue indices in sparse mode). A star is kept if it has at least one survey and at least one valid label (any parameter, any label source); per-parameter masks handle the rest. Splits are stored as row indices.

Functions

complete_config_from_catalogue(config_dict)

Fill in a raw config's data-dependent settings from its catalogue.

load_training_data(config[, loader, verbose])

Load the data an experiment config describes.

Classes

DataSplit(train, val, test)

Train/validation/test row indices.

TrainingData(kind, label_sources, valid_indices)

Experiment data loaded according to a config.

class dorothy.data.training_data.DataSplit(train, val, test)[source]

Bases: object

Train/validation/test row indices.

train

Training row indices.

val

Validation row indices.

test

Test row indices.

save(checkpoint_path, data)[source]

Save the split next to a checkpoint (data_split.pkl).

Parameters:
  • checkpoint_path (str | Path) – Checkpoint directory.

  • data (TrainingData) – The data the split indexes; its layout and valid rows are saved with the split.

Returns:

Path of the written file.

Return type:

Path

classmethod load(checkpoint_path, data)[source]

Load a saved split.

Split files without index_space: rows (single-survey and dense multi-label layouts) hold indices into the filtered arrays. They are mapped to rows with TrainingData.legacy_filtered_rows().

Parameters:
  • checkpoint_path (str | Path) – Checkpoint directory holding data_split.pkl.

  • data (TrainingData) – The data the split indexes.

Returns:

The split, as row indices.

Raises:
  • FileNotFoundError – If the directory has no data_split.pkl.

  • ValueError – If the split was saved for one survey and one label source, and data has another layout.

Return type:

DataSplit

class dorothy.data.training_data.TrainingData(kind, label_sources, valid_indices, X=None, y=None, X_dict=None, y_dict=None, has_data=None, has_labels=None, sparse=None)[source]

Bases: object

Experiment data loaded according to a config.

Only the fields for kind are set (see module docstring).

kind

Data layout.

label_sources

Label source names, in config order.

valid_indices

Rows kept for training/evaluation.

X, y

Single-survey spectra (N, 3, W) and labels (N, 3, P).

X_dict, y_dict, has_data, has_labels

Dense multi-survey arrays keyed by survey (X_dict, has_data) or label source (y_dict, has_labels).

sparse

Sparse multi-survey data.

property n_rows: int

Number of loaded rows (stars).

labels(source=None)[source]

Physical-unit labels for a label source.

Parameters:

source (str | None) – Label source name. None, or a name the sparse data does not hold separately, gives the primary labels.

Returns:

Array of shape (n_rows, 3, n_params) = [values, errors, mask].

Return type:

ndarray

split(train_ratio, val_ratio, seed)[source]

Randomly split the valid rows into train/validation/test.

Permutes the valid rows with numpy.random.RandomState(seed). The first int(n * train_ratio) go to train, the next int(n * val_ratio) to validation and the rest to test, where n is the number of valid rows.

Parameters:
  • train_ratio (float) – Fraction of valid rows for training.

  • val_ratio (float) – Fraction of valid rows for validation.

  • seed (int) – Seed of the permutation.

Returns:

DataSplit of row indices.

Return type:

DataSplit

configure_model(config)[source]

Set the MLP input features of a single-survey config, in place.

Parameters:

config (ExperimentConfig) – Experiment config, modified in place. Other layouts and configs without model are left unchanged.

fit(trainer, split)[source]

Train with the fit method matching this data layout.

Parameters:
  • trainer (Trainer) – The trainer.

  • split (DataSplit) – Row indices; only train and val are used.

Returns:

The training history.

Return type:

TrainingHistory

batch_source(rows, labels, device, survey_key=None)[source]

Batch source over the given rows, for evaluation.

Parameters:
  • rows (np.ndarray) – Row indices to iterate over.

  • labels (dict[str, np.ndarray]) – Label source name -> labels (n_rows, 3, P), e.g. normalized.

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

  • survey_key (str | None) – Input key for single-survey data (the survey name for a MultiHeadMLP, otherwise the default key).

Returns:

A sparse or tensor batch source over rows, with batches(size).

legacy_filtered_rows()[source]

Rows kept by the filter that old split files indexed into.

Returns:

Row indices. Single-survey – rows with a valid label. Dense multi-label: rows with a survey and a positive label error in any source. Sparse: valid_indices.

Return type:

ndarray

dorothy.data.training_data.complete_config_from_catalogue(config_dict, verbose=False)[source]

Fill in a raw config’s data-dependent settings from its catalogue.

Sets (from the catalogue, the authoritative source):

  • multi_head_model.survey_wavelengths for multi-head / multi-survey configs, and multi_head_model.label_sources when several label sources are configured;

  • data.parameter_names (an explicit value must match the catalogue).

Parameters:
  • config_dict (dict) – Raw config mapping (e.g. parsed YAML); not modified.

  • verbose (bool) – Print what was filled in.

Returns:

A completed copy of the config mapping, ready for ExperimentConfig.

Raises:
  • FileNotFoundError – If the catalogue does not exist.

  • ValueError – If data.catalogue_path is missing, the catalogue cannot provide the settings, or explicit parameter names disagree with it.

Return type:

dict

dorothy.data.training_data.load_training_data(config, loader=None, verbose=True)[source]

Load the data an experiment config describes.

Parameters:
  • config (ExperimentConfig) – Experiment configuration.

  • loader (CatalogueLoader | None) – Catalogue loader (created from data.catalogue_path if None).

  • verbose (bool) – Print a short summary.

Returns:

TrainingData with the valid rows identified.

Raises:
  • FileNotFoundError – If no loader is given and the catalogue does not exist.

  • ValueError – If the catalogue’s parameters don’t match the config, or a label source to duplicate from is not loaded.

Return type:

TrainingData