Reference / API reference / dorothy / dorothy.config / dorothy.config.schema

dorothy.config.schema

Pydantic configuration schemas for DOROTHY experiments.

This module defines the configuration structure for all aspects of a DOROTHY experiment, including data loading, model architecture, training parameters, and optional masking strategies.

All configurations use Pydantic v2 for validation and serialization.

Functions

default_parameter_names(n_parameters)	Return the first n_parameters standard names, then param_<i>.
parameter_display_name(name)	Return a plot label for a parameter name.

Classes

ActivationType(*values)	Supported activation function types.
CombinationMode(*values)	How to combine multi-survey embeddings.
DataConfig(*, catalogue_path[, surveys, ...])	Configuration for data loading and preprocessing.
ExperimentConfig(*, name[, description, ...])	Top-level configuration for a DOROTHY experiment.
InputMaskingConfig(*[, enabled, ...])	Configuration for dynamic input (spectrum) masking during training.
LabelMaskingConfig(*[, enabled, ...])	Configuration for dynamic label masking during training.
LossType(*values)	Supported loss function types.
ModelConfig(*[, hidden_layers, ...])	Configuration for the neural network architecture.
MultiHeadModelConfig(*, survey_wavelengths)	Configuration for multi-head MLP architecture.
NormalizationType(*values)	Supported normalization layer types.
OptimizerConfig(*[, type, weight_decay, ...])	Configuration for optimizers.
OptimizerType(*values)	Supported optimizer types.
SchedulerConfig(*[, type, base_lr, max_lr, ...])	Configuration for learning rate schedulers.
SchedulerType(*values)	Supported learning rate scheduler types.
TrainingConfig(*[, epochs, batch_size, ...])	Configuration for the training process.

class dorothy.config.schema.NormalizationType(*values)

Bases: str, Enum

Supported normalization layer types.

class dorothy.config.schema.ActivationType(*values)

Bases: str, Enum

Supported activation function types.

class dorothy.config.schema.LossType(*values)

Bases: str, Enum

Supported loss function types.

class dorothy.config.schema.OptimizerType(*values)

Bases: str, Enum

Supported optimizer types.

class dorothy.config.schema.SchedulerType(*values)

Bases: str, Enum

Supported learning rate scheduler types.

dorothy.config.schema.default_parameter_names(n_parameters)

Return the first n_parameters standard names, then param_<i>.

Parameters:

n_parameters (int) – Number of model output parameters.

Returns:

List of n_parameters parameter names.

Return type:

list[str]

dorothy.config.schema.parameter_display_name(name)

Return a plot label for a parameter name.

Abundance names of the form x_h / x_fe become [X/H] / [X/Fe] (e.g. mg_fe -> [Mg/Fe], fe_h -> [Fe/H]). Unrecognized names are returned unchanged.

Parameters:

name (str) – Parameter name as stored in the catalogue.

Returns:

Human-readable (matplotlib mathtext) label.

Return type:

str

class dorothy.config.schema.DataConfig(*, catalogue_path, surveys=['boss'], label_sources=['apogee'], train_ratio=0.7, val_ratio=0.2, max_flag_bits=0, duplicate_labels=None, use_dense_loading=False, parameter_names=None, normalize_spectra=True, precision_encoding=True)

Bases: _ConfigModel

Configuration for data loading and preprocessing.

Loads data from the super-catalogue HDF5 file which contains multiple surveys with cross-matched labels from APOGEE and/or GALAH.

Fields

Name	Type	Default	Description
catalogue_path	Path	required	Path to HDF5 super-catalogue
surveys	list	['boss']	List of input surveys to load (e.g., [‘boss’, ‘lamost_lrs’]). Several surveys give a merged dataset (outer join).
label_sources	list	['apogee']	List of label sources for training targets (e.g., [‘apogee’, ‘galah’])
train_ratio	float	0.7	Training data fraction
val_ratio	float	0.2	Validation data fraction
max_flag_bits	int	0	Maximum allowed flag bits (0=highest quality, -1=allow all with per-param masking)
duplicate_labels	dict[str, str] | None	None	Map target label source to source, to test multi-labelset training when a label source does not exist (e.g., {‘galah’: ‘apogee’} copies the APOGEE labels as GALAH labels)
use_dense_loading	bool	False	With several label sources, store spectra as dense arrays instead of sparse storage (more memory). No effect with one label source.
parameter_names	list[str] | None	None	Label parameter names in catalogue order. dorothy train fills them from the catalogue; when None, the first n STELLAR_PARAMETERS (n from the model section) are used, then param_<i>.
normalize_spectra	bool	True	Per-spectrum median/IQR flux normalization, with ivar rescaled to match. Training and inference must agree; the checkpoint config stores it.
precision_encoding	bool	True	Replace the sigma channel of the model inputs with log(1 + 1/sigma) at valid pixels and 0 at masked pixels, after normalization (encode_precision). False keeps sigma and exists only to load checkpoints trained with it; training rejects it.

property test_ratio: float

Compute test ratio from train and validation ratios.

property is_multi_survey: bool

Whether loading from multiple input surveys.

property is_multi_label: bool

Whether training with multiple label sources.

classmethod accept_singular_keys(values)

Accept survey / label_source as aliases of the list fields.

A string becomes a one-element list.

Raises:

ValueError – If a singular key and its plural are both given.

classmethod validate_parameter_names(v)

Ensure parameter names are unique.

validate_ratios()

Ensure train + val ratios don’t exceed 1.0.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.ModelConfig(*, hidden_layers=[5000, 2000, 1000, 500, 200, 100], normalization=NormalizationType.LAYERNORM, activation=ActivationType.GELU, dropout=0.0, input_features=15300, output_features=22, label_aware_uncertainty=True)

Bases: _ConfigModel

Configuration for the neural network architecture.

The default architecture follows the standard DOROTHY MLP: input_features -> 5000 -> 2000 -> 1000 -> 500 -> 200 -> 100 -> output_features (15300 and 22 by default).

Fields

Name	Type	Default	Description
hidden_layers	list	[5000, 2000, 1000, 500, 200, 100]	Hidden layer sizes
normalization	NormalizationType	<NormalizationType.LAYERNORM: 'layernorm'>	Normalization layer type
activation	ActivationType	<ActivationType.GELU: 'gelu'>	Activation function
dropout	float	0.0	Dropout probability
input_features	int	15300	Number of input features (channels * wavelength_bins)
output_features	int	22	Number of outputs (n_params * 2 for mean + uncertainty)
label_aware_uncertainty	bool	True	Predict ln s with an error decoder, parallel to the label decoder, from the head’s input and the reported label errors, or a learned missing vector, so that s includes the label noise (see LabelAwareOutput). False builds the output layers of checkpoints trained without it and exists only to load them; training rejects it.

classmethod validate_hidden_layers(v)

Ensure all hidden layer sizes are positive.

property n_parameters: int

Number of stellar parameters being predicted.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.CombinationMode(*values)

Bases: str, Enum

How to combine multi-survey embeddings.

class dorothy.config.schema.MultiHeadModelConfig(*, survey_wavelengths, n_parameters=11, latent_dim=256, encoder_hidden=[1024, 512], trunk_hidden=[512, 256], output_hidden=[64], combination_mode=CombinationMode.CONCAT, label_sources=None, normalization=NormalizationType.LAYERNORM, activation=ActivationType.GELU, dropout=0.0, head_dropout=<factory>, label_aware_uncertainty=True, residual_stream=True)

Bases: _ConfigModel

Configuration for multi-head MLP architecture.

Used when training on multiple surveys with different wavelength grids. Each survey has its own encoder, feeding into a shared trunk. Optionally supports multiple output heads for different label sources.

Examples

>>> config = MultiHeadModelConfig(
...     survey_wavelengths={"boss": 4506, "lamost_lrs": 3700},
...     n_parameters=11,
...     latent_dim=256,
...     label_sources=["apogee", "galah"],  # Multi-output heads
... )

Fields

Name	Type	Default	Description
survey_wavelengths	dict	required	Dict mapping survey names to wavelength bin counts
n_parameters	int	11	Number of stellar parameters to predict
latent_dim	int	256	Output dimension of each survey encoder
encoder_hidden	list	[1024, 512]	Hidden layer sizes for survey encoders
trunk_hidden	list	[512, 256]	Hidden layer sizes for shared trunk
output_hidden	list	[64]	Hidden layer sizes of each output head’s label and error decoders
combination_mode	CombinationMode	<CombinationMode.CONCAT: 'concat'>	How to combine multi-survey embeddings
label_sources	list[str] | None	None	Label sources of the output heads, one head each (e.g., [‘apogee’, ‘galah’]). None or an empty list gives one head, named after the first of data.label_sources.
normalization	NormalizationType	<NormalizationType.LAYERNORM: 'layernorm'>	Normalization layer type
activation	ActivationType	<ActivationType.GELU: 'gelu'>	Activation function
dropout	float	0.0	Dropout probability
head_dropout	dict	{}	Output head dropout per label source, e.g. {“galah”: 0.2}; sources not listed use dropout
label_aware_uncertainty	bool	True	Predict ln s with an error decoder, parallel to the label decoder, from the head’s input and the reported label errors, or a learned missing vector, so that s includes the label noise (see LabelAwareOutput). False builds the output layers of checkpoints trained without it and exists only to load them; training rejects it.
residual_stream	bool	True	Build the trunk and the output heads’ hidden layers as a residual stream of width latent_dim, with one residual block per trunk_hidden and output_hidden size; False builds plain blocks. Checkpoints saved before this option load with False.

classmethod validate_head_dropout(v)

Ensure per-head dropout probabilities are in [0, 1).

property is_multi_label: bool

Whether using multiple output heads for different label sources.

classmethod validate_survey_wavelengths(v)

Ensure survey wavelengths are valid.

property output_features: int

Number of output features (2 * n_parameters).

property survey_names: list[str]

List of survey names.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.SchedulerConfig(*, type=SchedulerType.ONE_CYCLE, base_lr=1e-06, max_lr=0.001, step_size_up=250, step_size_down=650, gamma=0.9995, patience=10, factor=0.5, pct_start=0.3, div_factor=25.0, final_div_factor=10000.0, anneal_strategy='cos')

Bases: _ConfigModel

Configuration for learning rate schedulers.

Fields

Name	Type	Default	Description
type	SchedulerType	<SchedulerType.ONE_CYCLE: 'one_cycle'>	Scheduler type
base_lr	float	1e-06	Minimum learning rate (cyclic)
max_lr	float	0.001	Learning rate: the peak of one_cycle and cyclic, the initial rate of reduce_on_plateau and cosine_annealing, and the constant rate of none
step_size_up	int	250	Steps to rise from base_lr to max_lr (cyclic)
step_size_down	int	650	Steps to fall from max_lr to base_lr (cyclic)
gamma	float	0.9995	Multiplicative factor of the exp_range mode (cyclic)
patience	int	10	Epochs to wait before reducing the LR (reduce_on_plateau)
factor	float	0.5	Factor to reduce the LR by (reduce_on_plateau)
pct_start	float	0.3	Fraction of the cycle spent increasing the LR (one_cycle)
div_factor	float	25.0	Initial LR = max_lr / div_factor (one_cycle)
final_div_factor	float	10000.0	Final LR = initial LR / final_div_factor (one_cycle)
anneal_strategy	Literal	'cos'	Annealing strategy (one_cycle)

validate_lr_range()

Ensure base_lr <= max_lr for the cyclic schedule, the one that uses both.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.OptimizerConfig(*, type=OptimizerType.ADAMW, weight_decay=0.01, betas=(0.9, 0.999), momentum=0.9, nesterov=False)

Bases: _ConfigModel

Configuration for optimizers.

Fields

Name	Type	Default	Description
type	OptimizerType	<OptimizerType.ADAMW: 'adamw'>	Optimizer type
weight_decay	float	0.01	Weight decay of the weight matrices. Biases and normalization weights are not decayed.
betas	tuple	(0.9, 0.999)	Coefficients for running averages (Adam/AdamW). The one_cycle scheduler replaces the first: it cycles it between 0.85 and 0.95.
momentum	float	0.9	Momentum factor (SGD only). The one_cycle scheduler replaces it: it cycles it between 0.85 and 0.95.
nesterov	bool	False	Use Nesterov momentum (SGD only)

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.TrainingConfig(*, epochs=300, batch_size=1024, loss=LossType.HETEROSCEDASTIC, optimizer=<factory>, scheduler=<factory>, gradient_clip=10.0, save_every=0, early_stopping_patience=0, scatter_floor=0.01, nll_warmup_fraction=0.1, label_error_dropout=0.5, calibrate_uncertainties=True, label_source_weighting='per_label', label_source_weights=<factory>, amp=False, amp_dtype='bfloat16', ema_decay=None, ema_cpu_every=None)

Bases: _ConfigModel

Configuration for the training process.

Fields

Name	Type	Default	Description
epochs	int	300	Number of epochs
batch_size	int	1024	Batch size
loss	LossType	<LossType.HETEROSCEDASTIC: 'heteroscedastic'>	Loss function
optimizer	OptimizerConfig	OptimizerConfig(type=<OptimizerType.ADAMW: 'adamw'>, weight_decay=0.01, betas=(0.9, 0.999), momentum=0.9, nesterov=False)	Optimizer configuration
scheduler	SchedulerConfig	SchedulerConfig(type=<SchedulerType.ONE_CYCLE: 'one_cycle'>, base_lr=1e-06, max_lr=0.001, step_size_up=250, step_size_down=650, gamma=0.9995, patience=10, factor=0.5, pct_start=0.3, div_factor=25.0, final_div_factor=10000.0, anneal_strategy='cos')	LR scheduler
gradient_clip	float	10.0	Max gradient norm (0=disabled)
save_every	int	0	Save a checkpoint every N epochs (0 = only the final one)
early_stopping_patience	int	0	Epochs without a lower validation loss before training stops, counted after the β-NLL warm-up; 0 turns early stopping off
scatter_floor	float	0.01	Minimum scatter floor for heteroscedastic loss
nll_warmup_fraction	float	0.1	Fraction of the optimizer steps of the run over which β of the β-NLL loss falls linearly from 1 (mean-squared-error gradients for the means) to 0 (the plain negative log-likelihood); 0 turns this warm-up off. Applies to the heteroscedastic loss only.
label_error_dropout	float	0.5	Fraction of training stars whose label errors are replaced by the learned missing vector, drawn per batch and output head. Validation, evaluation and prediction always use the missing vector, so the fraction must be above 0.
calibrate_uncertainties	bool	True	After training, fit one factor per output head and parameter that scales the predicted scatter so that the validation z-scores of the best epoch have unit robust spread, and save the factors in calibration.json. loss: mse has no scatter, so it saves no factors.
label_source_weighting	Literal	'per_label'	How the losses of the label sources (output heads) combine, for training and for selecting the best epoch. per_label gives every valid label equal weight; per_source averages each head’s mean loss, so a smaller label source gets as much total weight as a larger one.
label_source_weights	dict	{}	Relative loss weight per label source, e.g. {“galah”: 0.5}; sources not listed have weight 1.0
amp	bool	False	Run the training forward pass in amp_dtype on CUDA. CPU and MPS train in float32.
amp_dtype	Literal	'bfloat16'	Precision of the forward pass with amp. bfloat16 needs a GPU with compute capability 8.0 or higher; float16 uses a gradient scaler.
ema_decay	float | None	None	Decay per optimizer step of the exponential moving average of the weights, which validation, the best epoch and the saved weights then use (null: no average)
ema_cpu_every	int | None	None	Keep the weight average in CPU memory and update it every this many optimizer steps (null: on the training device, every step)

classmethod validate_label_source_weights(v)

Ensure label source weights are non-negative.

validate_ema()

Reject a weight-average option without a weight average.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.LabelMaskingConfig(*, enabled=False, p_labelset_min=0.3, p_labelset_max=1.0, p_label_min=0.3, p_label_max=1.0)

Bases: _ConfigModel

Configuration for dynamic label masking during training.

Label masking is a training-time augmentation that randomly masks labels to improve model robustness when not all labels are available. This uses a hierarchical scheme:

1. Labelset level: Which label sources (apogee, galah, etc.) to include

2. Label level: Which parameters (teff, logg, feh, etc.) within each set

All probabilities are sampled fresh per batch from uniform ranges to ensure the model sees diverse masking conditions during training.

A guaranteed keeper mechanism ensures at least one labelset and one label per labelset are always kept, preventing complete masking.

Fields

Name	Type	Default	Description
enabled	bool	False	Enable label masking augmentation
p_labelset_min	float	0.3	Minimum probability of keeping each labelset (sampled per batch)
p_labelset_max	float	1.0	Maximum probability of keeping each labelset (sampled per batch)
p_label_min	float	0.3	Minimum probability of keeping each label within kept labelsets (sampled per batch)
p_label_max	float	1.0	Maximum probability of keeping each label within kept labelsets (sampled per batch)

validate_ranges()

Validate probability ranges.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.InputMaskingConfig(*, enabled=False, p_survey_min=0.3, p_survey_max=1.0, f_min_override=None, f_max=0.5, p_block_min=0.3, p_block_max=1.0)

Bases: _ConfigModel

Configuration for dynamic input (spectrum) masking during training.

Input masking is a training-time augmentation that randomly masks portions of input spectra to improve model robustness. This uses a hierarchical scheme:

1. Survey level: Which surveys to include (for multi-survey training)

2. Block level: Which wavelength blocks within each kept survey

Block sizes are sampled log-uniformly to explore all scales from single pixels up to large spectral regions. All probabilities are sampled fresh per batch from uniform ranges.

A guaranteed keeper mechanism ensures at least one survey and one block per survey are always kept, preventing complete masking.

Fields

Name	Type	Default	Description
enabled	bool	False	Enable input masking augmentation
p_survey_min	float	0.3	Minimum probability of keeping each survey (sampled per batch)
p_survey_max	float	1.0	Maximum probability of keeping each survey (sampled per batch)
f_min_override	float | None	None	Override for minimum block size fraction. If None, defaults to 1/N_wavelengths (single pixel).
f_max	float	0.5	Maximum block size as fraction of spectrum
p_block_min	float	0.3	Minimum probability of keeping each block (sampled per batch)
p_block_max	float	1.0	Maximum probability of keeping each block (sampled per batch)

validate_ranges()

Validate probability and fraction ranges.

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class dorothy.config.schema.ExperimentConfig(*, name, description='', data, model=None, multi_head_model=None, training=<factory>, label_masking=<factory>, input_masking=<factory>, output_dir=PosixPath('outputs'), seed=42, device='auto')

Bases: _ConfigModel

Top-level configuration for a DOROTHY experiment.

This is the main configuration class that combines all sub-configurations and provides the complete specification for training a model.

For single-survey training, use the model config (MLP). For multi-survey training, use the multi_head_model config (MultiHeadMLP).

Fields

Name	Type	Default	Description
name	str	required	Experiment name
description	str	''	Experiment description
data	DataConfig	required	Data configuration
model	dorothy.config.schema.ModelConfig | None	None	Standard MLP configuration (single-survey)
multi_head_model	dorothy.config.schema.MultiHeadModelConfig | None	None	Multi-head MLP configuration (multi-survey)
training	TrainingConfig	TrainingConfig(epochs=300, batch_size=1024, loss=<LossType.HETEROSCEDASTIC: 'heteroscedastic'>, optimizer=OptimizerConfig(type=<OptimizerType.ADAMW: 'adamw'>, weight_decay=0.01, betas=(0.9, 0.999), momentum=0.9, nesterov=False), scheduler=SchedulerConfig(type=<SchedulerType.ONE_CYCLE: 'one_cycle'>, base_lr=1e-06, max_lr=0.001, step_size_up=250, step_size_down=650, gamma=0.9995, patience=10, factor=0.5, pct_start=0.3, div_factor=25.0, final_div_factor=10000.0, anneal_strategy='cos'), gradient_clip=10.0, save_every=0, early_stopping_patience=0, scatter_floor=0.01, nll_warmup_fraction=0.1, label_error_dropout=0.5, calibrate_uncertainties=True, label_source_weighting='per_label', label_source_weights={}, amp=False, amp_dtype='bfloat16', ema_decay=None, ema_cpu_every=None)	Training configuration
label_masking	LabelMaskingConfig	LabelMaskingConfig(enabled=False, p_labelset_min=0.3, p_labelset_max=1.0, p_label_min=0.3, p_label_max=1.0)	Dynamic label masking configuration
input_masking	InputMaskingConfig	InputMaskingConfig(enabled=False, p_survey_min=0.3, p_survey_max=1.0, f_min_override=None, f_max=0.5, p_block_min=0.3, p_block_max=1.0)	Dynamic input masking configuration
output_dir	Path	PosixPath('outputs')	Directory for outputs (checkpoints, logs)
seed	int	42	Random seed
device	Literal	'auto'	Compute device (mps for Apple Silicon)

property is_multi_head: bool

Whether using multi-head architecture.

property parameter_names: list[str]

Names of the predicted parameters, in model output order.

Uses data.parameter_names when set (populated from the catalogue), otherwise the first n STELLAR_PARAMETERS, followed by param_<i> for each output beyond that list.

validate_model_config()

Ensure model configuration is consistent with data configuration.

get_output_path()

Get the output directory of this experiment.

Returns:

output_dir / name.

Return type:

Path

get_checkpoint_path(epoch=None)

Get the directory of a checkpoint.

Parameters:

epoch (int | None) – Epoch number, or None for the final checkpoint.

Returns:

<output path>/epoch_<epoch>, or <output path>/<name>_final when epoch is None.

Return type:

Path

save_yaml(path)

Save configuration to a YAML file.

Parameters:

path (str | Path) – Path to save the YAML file.

classmethod load_checkpoint_config(path)

Load the config.yaml saved alongside a trained checkpoint.

Unlike load_yaml(), this accepts the Python-specific tags (e.g. tuples) written by save_yaml(), and fills in defaults that keep older checkpoints reproducible. A config without data.normalize_spectra was trained on un-normalized catalogue spectra, so the flag defaults to False, except for configs from the dorothy_update branch, recognizable by a multi_head_model.head_dropout entry or 20 predicted parameters, whose sparse and single-survey loaders already normalized spectra. A warning is issued either way. A model section without label_aware_uncertainty, a multi-head model section without residual_stream and a data section without precision_encoding get False for that flag. A training.learning_rate becomes scheduler.max_lr for the schedules it set (none, reduce_on_plateau, cosine_annealing). A model section with several label sources keeps only the first, the one it trained on. Settings this version no longer has are ignored with a warning.

Parameters:

path (str | Path) – Path to the checkpoint’s config.yaml.

Returns:

ExperimentConfig instance.

Raises:

pydantic.ValidationError – If the settings are invalid.

Return type:

ExperimentConfig

classmethod load_yaml(path)

Load configuration from a YAML file.

Parameters:

path (str | Path) – Path to the YAML file.

Returns:

ExperimentConfig instance.

Return type:

ExperimentConfig

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].