Reference / API reference / dorothy / dorothy.visualization / dorothy.visualization.training_plots

dorothy.visualization.training_plots

Plots of a training run, drawn from a TrainingHistory.

Each plot_* function writes one PNG into a directory and returns its path, or None when matplotlib is not installed or the history lacks the data. generate_training_report calls them all.

Functions

generate_training_report(history, output_dir)	Write every training plot that the history supports.
plot_grokking_metrics(history, output_path)	Plot the gradient norm, weight norms and weight updates per epoch.
plot_loss_components(history, output_path)	Plot the mean and scatter loss components per epoch.
plot_loss_curves(history, output_path[, title])	Plot the loss curves and the learning rate schedule.
plot_metrics_evolution(history, output_path)	Plot the validation RMSE and z-score scatter per epoch.
plot_per_labelset_loss_curves(history, ...)	Plot the validation loss of each label source, absolute and relative.
plot_per_labelset_metrics(history, output_path)	Plot the validation metrics of each label source at the best epoch.
plot_per_labelset_metrics_evolution(history, ...)	Plot the validation RMSE of each label source per epoch.
plot_per_survey_loss_curves(history, output_path)	Plot the validation loss of each survey, absolute and relative.
plot_per_survey_metrics(history, output_path)	Plot the validation metrics of each survey at the best epoch.
plot_per_survey_metrics_evolution(history, ...)	Plot the validation RMSE of each survey per epoch.
plot_val_metrics(history, output_path)	Plot the validation RMSE, robust scatter, bias and MAE per parameter.
plot_zscore_calibration(history, output_path)	Plot the median and robust scatter of the z-scores per parameter.

dorothy.visualization.training_plots.plot_loss_curves(history, output_path, title='Training Progress')

Plot the loss curves and the learning rate schedule.

The left panel shows the training and validation loss, the β-NLL warm-up and the best epoch. The right panel shows the learning rate per epoch, or per step when the history has one value per step.

Parameters:

• history (TrainingHistory) – Training history with train_losses.

• output_path (str | Path) – Directory of the plot, written as loss_curves.png.

• title (str) – Figure title.

Returns:

Path of the plot, or None if matplotlib is missing or train_losses is empty.

Return type:

Path | None

dorothy.visualization.training_plots.plot_loss_components(history, output_path)

Plot the mean and scatter loss components per epoch.

Shows the first 6 parameters.

Parameters:

• history (TrainingHistory) – Training history with val_loss_breakdown.

• output_path (str | Path) – Directory of the plot, written as loss_components.png.

Returns:

Path of the plot, or None if matplotlib is missing or the breakdown is empty or not of shape (n_epochs, n_params).

Return type:

Path | None

dorothy.visualization.training_plots.plot_grokking_metrics(history, output_path)

Plot the gradient norm, weight norms and weight updates per epoch.

The weight panels show the mean over the layers of each component (encoders per survey, trunk, output heads, plain layers). They appear only when the history has weight_norms or weight_updates.

Parameters:

• history (TrainingHistory) – Training history with grad_norms.

• output_path (str | Path) – Directory of the plot, written as grokking_metrics.png.

Returns:

Path of the plot, or None if matplotlib is missing or grad_norms is empty.

Return type:

Path | None

dorothy.visualization.training_plots.plot_val_metrics(history, output_path)

Plot the validation RMSE, robust scatter, bias and MAE per parameter.

The metrics are those of the best epoch, whose weights best_model.pth holds, in normalized units: (value - median) / IQR, with Teff in log10. Without robust_scatter the scatter panel shows sd. Bias is predicted minus true.

Parameters:

• history (TrainingHistory) – Training history with val_metrics.

• output_path (str | Path) – Directory of the plot, written as val_metrics.png.

Returns:

Path of the plot, or None if matplotlib is missing or rmse is empty.

Return type:

Path | None

dorothy.visualization.training_plots.plot_zscore_calibration(history, output_path)

Plot the median and robust scatter of the z-scores per parameter.

Calibrated uncertainties give z-scores with median 0 and robust scatter 1. The values are those of the best epoch, whose weights best_model.pth holds.

Parameters:

• history (TrainingHistory) – Training history with z_median and z_robust_scatter in val_metrics.

• output_path (str | Path) – Directory of the plot, written as zscore_calibration.png.

Returns:

Path of the plot, or None if matplotlib is missing or either metric is empty.

Return type:

Path | None

dorothy.visualization.training_plots.plot_metrics_evolution(history, output_path)

Plot the validation RMSE and z-score scatter per epoch.

Shows the first 3 parameters. Without z-score scatter for every epoch, the right panel shows the validation loss.

Parameters:

• history (TrainingHistory) – Training history with val_metrics.

• output_path (str | Path) – Directory of the plot, written as metrics_evolution.png.

Returns:

Path of the plot, or None if matplotlib is missing or rmse has fewer than 2 epochs.

Return type:

Path | None

dorothy.visualization.training_plots.plot_per_survey_loss_curves(history, output_path)

Plot the validation loss of each survey, absolute and relative.

The left panel shows the loss of each survey and the overall loss. The right panel shows the survey loss divided by the overall loss.

Parameters:

• history (TrainingHistory) – Training history with per_survey_val_losses.

• output_path (str | Path) – Directory of the plot, written as per_survey_loss_curves.png.

Returns:

Path of the plot, or None if matplotlib is missing or the history has no per-survey losses.

Return type:

Path | None

dorothy.visualization.training_plots.plot_per_survey_metrics(history, output_path)

Plot the validation metrics of each survey at the best epoch.

The 2x2 grid of grouped bars shows the RMSE, robust scatter, bias and MAE per parameter, as in plot_val_metrics, for the first 6 parameters.

Parameters:

• history (TrainingHistory) – Training history with per_survey_val_metrics.

• output_path (str | Path) – Directory of the plot, written as per_survey_metrics.png.

Returns:

Path of the plot, or None if matplotlib is missing or the first survey has no rmse.

Return type:

Path | None

dorothy.visualization.training_plots.plot_per_survey_metrics_evolution(history, output_path)

Plot the validation RMSE of each survey per epoch.

One panel per parameter, for the first 3 parameters.

Parameters:

• history (TrainingHistory) – Training history with per_survey_val_metrics.

• output_path (str | Path) – Directory of the plot, written as per_survey_metrics_evolution.png.

Returns:

Path of the plot, or None if matplotlib is missing or the first survey has fewer than 2 epochs of rmse.

Return type:

Path | None

dorothy.visualization.training_plots.plot_per_labelset_loss_curves(history, output_path)

Plot the validation loss of each label source, absolute and relative.

The left panel shows the loss of each output head and the overall loss. The right panel shows the head loss divided by the overall loss.

Parameters:

• history (TrainingHistory) – Training history with per_labelset_val_losses.

• output_path (str | Path) – Directory of the plot, written as per_labelset_loss_curves.png.

Returns:

Path of the plot, or None if matplotlib is missing or the history has no per-label-source losses.

Return type:

Path | None

dorothy.visualization.training_plots.plot_per_labelset_metrics(history, output_path)

Plot the validation metrics of each label source at the best epoch.

The 2x2 grid of grouped bars shows the RMSE, robust scatter, bias and MAE per parameter, as in plot_val_metrics, for the first 6 parameters.

Parameters:

• history (TrainingHistory) – Training history with per_labelset_val_metrics.

• output_path (str | Path) – Directory of the plot, written as per_labelset_metrics.png.

Returns:

Path of the plot, or None if matplotlib is missing or the first label source has no rmse.

Return type:

Path | None

dorothy.visualization.training_plots.plot_per_labelset_metrics_evolution(history, output_path)

Plot the validation RMSE of each label source per epoch.

One panel per parameter, for the first 3 parameters.

Parameters:

• history (TrainingHistory) – Training history with per_labelset_val_metrics.

• output_path (str | Path) – Directory of the plot, written as per_labelset_metrics_evolution.png.

Returns:

Path of the plot, or None if matplotlib is missing or the first label source has fewer than 2 epochs of rmse.

Return type:

Path | None

dorothy.visualization.training_plots.generate_training_report(history, output_dir, experiment_name='training')

Write every training plot that the history supports.

The per-survey plots need survey_names and per_survey_val_losses. The per-label-source plots need label_source_names and per_labelset_val_losses. A plot that raises is skipped with a printed warning.

Parameters:

• history (TrainingHistory) – Training history.

• output_dir (str | Path) – Directory of the plots, created if missing.

• experiment_name (str) – Used in the title of the loss curves.

Returns:

Paths of the plots written.

Return type:

list[Path]