Skip to content

Add opt-in error-aware emulator training #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
botjean wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add opt-in error-aware emulator training #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions src/bayesian/data_IO.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1462,14 +1462,20 @@ def read_dict_from_h5(input_dir: Path, filename: str, verbose: bool = True) -> d


def predictions_matrix_from_h5(
output_dir, filename, validation_set=False, observable_filter: ObservableFilter | None = None
output_dir,
filename,
validation_set=False,
observable_filter: ObservableFilter | None = None,
value_key: str = "y",
):
"""
Initialize predictions from observables.h5 file into a single 2D array:

:param str output_dir: location of filename
:param str filename: h5 filename (typically 'observables.h5')
:param ObservableFilter observable_filter: (optional) filter to apply to the observables
:param str value_key: prediction field to stack (for example, ``y`` or
``y_err_stat``)
:return 2darray Y: matrix of predictions at all design points (design_point_index, observable_bins) i.e. (n_samples, n_features)
"""

Expand All @@ -1489,7 +1495,7 @@ def predictions_matrix_from_h5(
# (design_point_index, observable_bins) i.e. (n_samples, n_features)
length_of_Y = 0
for i, observable_label in enumerate(sorted_observable_list):
values = observables[prediction_label][observable_label]["y"].T
values = observables[prediction_label][observable_label][value_key].T
length_of_Y += values.shape[1]
logger.info(f"{observable_label} shape: {values.shape}, length: {length_of_Y=}")
if i == 0:
Expand Down
245 changes: 226 additions & 19 deletions src/bayesian/emulation/sk_learn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@

from __future__ import annotations

import hashlib
import json
import logging
from pathlib import Path
from typing import Any, ClassVar
Expand All @@ -34,6 +36,109 @@

# Name under which the module is registered.
_register_name = "sk_learn"
_CACHE_METADATA_VERSION = 1


def _training_metadata(
emulator_settings: SKLearnEmulatorSettings,
analysis_settings: analysis.AnalysisSettings,
Y: npt.NDArray[np.float64],
design: npt.NDArray[np.float64],
Y_err_stat: npt.NDArray[np.float64] | None = None,
) -> dict[str, Any]:
"""Describe the inputs and settings that make a cached GP reusable."""
observable_list = emulator_settings.settings.get("observable_list", [])
observable_names = [
item["observable"] if isinstance(item, dict) else item
for item in observable_list
]
parameterization = analysis_settings.raw_analysis_config["parameterization"][
analysis_settings.parameterization
]
settings = {
"backend": _register_name,
"n_pc": emulator_settings.n_pc,
"max_n_components_to_calculate": emulator_settings.max_n_components_to_calculate,
"active_kernels": emulator_settings.active_kernels,
"n_restarts": emulator_settings.n_restarts,
"alpha": emulator_settings.alpha,
"normalize_y": emulator_settings.normalize_y,
"random_state": emulator_settings.random_state,
"use_prediction_statistical_uncertainty": (
emulator_settings.use_prediction_statistical_uncertainty
),
"observable_names": observable_names,
"observable_exclude_list": emulator_settings.settings.get(
"observable_exclude_list", []
),
"parameterization": {
"name": analysis_settings.parameterization,
"parameter_names": parameterization.get("names"),
"min": parameterization["min"],
"max": parameterization["max"],
},
}
settings_json = json.dumps(
settings,
sort_keys=True,
separators=(",", ":"),
).encode()

arrays_hash = hashlib.sha256()
arrays = [("Y", Y), ("design", design)]
if Y_err_stat is not None:
arrays.append(("Y_err_stat", Y_err_stat))
for name, array in arrays:
canonical = np.ascontiguousarray(array, dtype="<f8")
arrays_hash.update(name.encode())
arrays_hash.update(str(canonical.shape).encode())
arrays_hash.update(canonical.tobytes())

return {
"version": _CACHE_METADATA_VERSION,
"settings": settings,
"settings_sha256": hashlib.sha256(settings_json).hexdigest(),
"training_arrays_sha256": arrays_hash.hexdigest(),
}


def _validate_cached_emulator(
cached_results: dict[str, Any],
expected_metadata: dict[str, Any],
output_filename: Path,
) -> None:
metadata = cached_results.get("training_metadata")
if metadata is None:
msg = (
f"Cached emulator {output_filename} has no training metadata, "
"so compatibility with the requested data and GP settings cannot "
"be verified. Set force_retrain: true once to replace this legacy "
"cache."
)
raise RuntimeError(msg)

if metadata != expected_metadata:
msg = (
f"Cached emulator {output_filename} was trained with different "
"settings or input arrays. Set force_retrain: true to replace it."
)
raise RuntimeError(msg)


def _build_gaussian_process(
kernel: Any,
emulator_settings: SKLearnEmulatorSettings,
*,
alpha: float | npt.NDArray[np.float64] | None = None,
) -> sklearn_gaussian_process.GaussianProcessRegressor:
return sklearn_gaussian_process.GaussianProcessRegressor(
kernel=kernel,
alpha=emulator_settings.alpha if alpha is None else alpha,
n_restarts_optimizer=emulator_settings.n_restarts,
normalize_y=emulator_settings.normalize_y,
random_state=emulator_settings.random_state,
copy_X_train=False,
)


def fit_emulator(
Expand All @@ -54,24 +159,65 @@ def fit_emulator(
emulator_settings=emulator_settings, analysis_settings=analysis_settings
)

# Load the exact training inputs before checking the cache so incompatible
# settings or data cannot silently reuse a fitted emulator.
Y = data_IO.predictions_matrix_from_h5(
output_dir=analysis_settings.output_dir,
filename=analysis_settings.io.observables_filename,
observable_filter=emulator_settings.base_settings.observable_filter,
)
design = data_IO.design_array_from_h5(
analysis_settings.output_dir,
filename=analysis_settings.io.observables_filename,
)
Y_err_stat = None
if emulator_settings.use_prediction_statistical_uncertainty:
Y_err_stat = data_IO.predictions_matrix_from_h5(
output_dir=analysis_settings.output_dir,
filename=analysis_settings.io.observables_filename,
observable_filter=emulator_settings.base_settings.observable_filter,
value_key="y_err_stat",
)
if Y_err_stat.shape != Y.shape:
raise ValueError(
"Prediction values and statistical uncertainties have "
f"different shapes: {Y.shape} != {Y_err_stat.shape}"
)
if not np.all(np.isfinite(Y_err_stat)) or np.any(Y_err_stat < 0):
raise ValueError(
"Prediction statistical uncertainties must be finite and "
"non-negative"
)
training_metadata = _training_metadata(
emulator_settings,
analysis_settings,
Y,
design,
Y_err_stat,
)

# Check if emulator already exists
if output_filename.exists():
if emulator_settings.force_retrain:
output_filename.unlink()
logger.info(f"Removed {output_filename}")
else:
cached_results = emulation_base.IO.read_emulator(
emulator_settings=emulator_settings,
analysis_settings=analysis_settings,
)
_validate_cached_emulator(
cached_results,
training_metadata,
output_filename,
)
logger.info(f"Emulators already exist: {output_filename} (to force retrain, set force_retrain: True)")
return {}

# Initialize predictions into a single 2D array: (design_point_index, observable_bins) i.e. (n_samples, n_features)
# A consistent order of observables is enforced internally in data_IO
# NOTE: One sample corresponds to one design point, while one feature is one bin of one observable
logger.info("Doing PCA...")
Y = data_IO.predictions_matrix_from_h5(
output_dir=analysis_settings.output_dir,
filename=analysis_settings.io.observables_filename,
observable_filter=emulator_settings.base_settings.observable_filter,
)

# Use sklearn to:
# - Center and scale each feature (and later invert)
Expand Down Expand Up @@ -118,6 +264,22 @@ def fit_emulator(
# Scale data and perform PCA
Y_pca = pca.fit_transform(scaler.fit_transform(Y))
Y_pca_truncated = Y_pca[:, : emulator_settings.n_pc] # Select PCs here
alpha_per_pc = _project_prediction_uncertainties_to_pc_space(
Y_err_stat=Y_err_stat,
scaler=scaler,
pca=pca,
Y_pca_truncated=Y_pca_truncated,
normalize_y=emulator_settings.normalize_y,
jitter=emulator_settings.alpha,
)
if alpha_per_pc is not None:
logger.info(
"Using projected model-statistical variances as heteroscedastic "
"GP alpha (min=%.3g, median=%.3g, max=%.3g)",
float(np.min(alpha_per_pc)),
float(np.median(alpha_per_pc)),
float(np.max(alpha_per_pc)),
)
# Invert PCA and undo the scaling
Y_reconstructed_truncated = Y_pca_truncated.dot(pca.components_[: emulator_settings.n_pc, :])
Y_reconstructed_truncated_unscaled = scaler.inverse_transform(Y_reconstructed_truncated)
Expand All @@ -126,11 +288,6 @@ def fit_emulator(
f" Variance explained by first {emulator_settings.n_pc} components: {np.sum(explained_variance_ratio[: emulator_settings.n_pc])}"
)

# Get design
design = data_IO.design_array_from_h5(
analysis_settings.output_dir, filename=analysis_settings.io.observables_filename
)

# Define GP kernel (covariance function)
min = np.array(analysis_settings.raw_analysis_config["parameterization"][analysis_settings.parameterization]["min"])
max = np.array(analysis_settings.raw_analysis_config["parameterization"][analysis_settings.parameterization]["max"])
Expand Down Expand Up @@ -173,15 +330,20 @@ def fit_emulator(
logger.info("")
logger.info("Fitting GPs...")
logger.info(f" The design has {design.shape[1]} parameters")
emulators = [
sklearn_gaussian_process.GaussianProcessRegressor(
kernel=kernel,
alpha=emulator_settings.alpha,
n_restarts_optimizer=emulator_settings.n_restarts,
copy_X_train=False,
).fit(design, y)
for y in Y_pca_truncated.T
]
emulators = []
for i_pc, y in enumerate(Y_pca_truncated.T):
alpha = (
emulator_settings.alpha
if alpha_per_pc is None
else alpha_per_pc[:, i_pc]
)
emulators.append(
_build_gaussian_process(
kernel,
emulator_settings,
alpha=alpha,
).fit(design, y)
)

# Print hyperparameters.
logger.info("")
Expand All @@ -199,11 +361,48 @@ def fit_emulator(
output_dict["PCA"]["Y_reconstructed_truncated_unscaled"] = Y_reconstructed_truncated_unscaled
output_dict["PCA"]["pca"] = pca
output_dict["PCA"]["scaler"] = scaler
output_dict["PCA"]["Y_err_stat"] = Y_err_stat
output_dict["PCA"]["alpha_per_pc"] = alpha_per_pc
output_dict["emulators"] = emulators
output_dict["training_metadata"] = training_metadata

return output_dict


def _project_prediction_uncertainties_to_pc_space(
*,
Y_err_stat: npt.NDArray[np.float64] | None,
scaler: sklearn_preprocessing.StandardScaler,
pca: sklearn_decomposition.PCA,
Y_pca_truncated: npt.NDArray[np.float64],
normalize_y: bool,
jitter: float,
) -> npt.NDArray[np.float64] | None:
"""Project diagonal feature-level simulation covariance into PC variances.

Off-diagonal PC covariance is discarded because each retained PC is fit by
an independent Gaussian process.
"""
if Y_err_stat is None:
return None

scaled_variance = (Y_err_stat / scaler.scale_[np.newaxis, :]) ** 2
retained_components_squared = (
pca.components_[: Y_pca_truncated.shape[1], :] ** 2
)
alpha_per_pc = scaled_variance @ retained_components_squared.T

# sklearn normalizes y before adding alpha, so alpha must use the same
# normalized target units.
if normalize_y:
pc_scale = np.std(Y_pca_truncated, axis=0)
pc_scale[pc_scale == 0] = 1.0
alpha_per_pc /= pc_scale[np.newaxis, :] ** 2

alpha_per_pc += jitter
return alpha_per_pc


def predict(
parameters: npt.NDArray[np.float64],
results: dict[str, Any],
Expand Down Expand Up @@ -402,6 +601,9 @@ class SKLearnEmulatorSettings:
settings: dict[str, Any]
# Additional name, for providing
additional_name: str = attrs.field(default="")
normalize_y: bool = attrs.field(default=False)
random_state: int | None = attrs.field(default=None)
use_prediction_statistical_uncertainty: bool = attrs.field(default=False)

def __attrs_post_init__(self):
"""
Expand Down Expand Up @@ -439,6 +641,11 @@ def from_config(cls, config: dict[str, Any]) -> SKLearnEmulatorSettings:
active_kernels={kernel_type: config["kernels"][kernel_type] for kernel_type in config["kernels"]["active"]},
n_restarts=config["GPR"]["n_restarts"],
alpha=config["GPR"]["alpha"],
normalize_y=config["GPR"].get("normalize_y", False),
random_state=config["GPR"].get("random_state"),
use_prediction_statistical_uncertainty=config["GPR"].get(
"use_prediction_statistical_uncertainty", False
),
settings=config,
)

Expand Down
24 changes: 24 additions & 0 deletions tests/test_data_IO.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,30 @@ def test_observable_matrix_round_trip(caplog: Any) -> None:
Y_round_trip = data_IO.observable_matrix_from_dict(Y_dict)
np.testing.assert_allclose(Y, Y_round_trip)


def test_predictions_matrix_selects_requested_value_key(monkeypatch: Any) -> None:
observable = "5020__PbPb__hadron__pt_ch_cms____0-5"
predictions = {
observable: {
"y": np.array([[0.4, 0.5], [0.6, 0.7]]),
"y_err_stat": np.array([[0.01, 0.02], [0.03, 0.04]]),
}
}
monkeypatch.setattr(
data_IO,
"read_dict_from_h5",
lambda *args, **kwargs: {"Prediction": predictions},
)

result = data_IO.predictions_matrix_from_h5(
".",
"unused.h5",
value_key="y_err_stat",
)

np.testing.assert_allclose(result, predictions[observable]["y_err_stat"].T)


@pytest.mark.parametrize(
"design_points_to_exclude",
[[17, 43, 203], []],
Expand Down
Loading