diff --git a/src/bayesian/data_IO.py b/src/bayesian/data_IO.py index 1e836ab..02063e5 100644 --- a/src/bayesian/data_IO.py +++ b/src/bayesian/data_IO.py @@ -1462,7 +1462,11 @@ 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: @@ -1470,6 +1474,8 @@ def predictions_matrix_from_h5( :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) """ @@ -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: diff --git a/src/bayesian/emulation/sk_learn.py b/src/bayesian/emulation/sk_learn.py index 89ca777..18ed2fe 100644 --- a/src/bayesian/emulation/sk_learn.py +++ b/src/bayesian/emulation/sk_learn.py @@ -15,6 +15,8 @@ from __future__ import annotations +import hashlib +import json import logging from pathlib import Path from typing import Any, ClassVar @@ -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=" 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( @@ -54,12 +159,58 @@ 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 {} @@ -67,11 +218,6 @@ def fit_emulator( # 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) @@ -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) @@ -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"]) @@ -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("") @@ -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], @@ -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): """ @@ -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, ) diff --git a/tests/test_data_IO.py b/tests/test_data_IO.py index 0f5b31e..0afce96 100644 --- a/tests/test_data_IO.py +++ b/tests/test_data_IO.py @@ -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], []], diff --git a/tests/test_emulation_methodology.py b/tests/test_emulation_methodology.py new file mode 100644 index 0000000..df3b4e8 --- /dev/null +++ b/tests/test_emulation_methodology.py @@ -0,0 +1,333 @@ +import pickle +from types import SimpleNamespace + +import numpy as np +import pytest +from silx.io.dictdump import dicttoh5 +from sklearn.decomposition import PCA +from sklearn.gaussian_process.kernels import Matern, WhiteKernel +from sklearn.preprocessing import StandardScaler + +from bayesian.emulation import base as emulation_base +from bayesian.emulation import sk_learn + + +def _emulator_config( + *, + normalize_y: bool | None = None, + use_prediction_statistical_uncertainty: bool | None = None, +) -> dict: + config = { + "force_retrain": False, + "emulator_package": "sk_learn", + "n_pc": 1, + "kernels": { + "active": ["matern"], + "matern": { + "nu": 1.5, + "length_scale_bounds_factor": [0.1, 10.0], + }, + }, + "GPR": { + "n_restarts": 0, + "alpha": 1.0e-8, + }, + "observable_list": ["test"], + } + if normalize_y is not None: + config["GPR"]["normalize_y"] = normalize_y + if use_prediction_statistical_uncertainty is not None: + config["GPR"]["use_prediction_statistical_uncertainty"] = ( + use_prediction_statistical_uncertainty + ) + return config + + +def _analysis_settings(tmp_path=None): + return SimpleNamespace( + output_dir=tmp_path, + io=SimpleNamespace(observables_filename="observables.h5"), + parameterization="unit", + raw_analysis_config={ + "parameterization": { + "unit": { + "names": ["x"], + "min": [0.0], + "max": [1.0], + } + } + }, + ) + + +def test_methodology_options_are_explicit_and_backward_compatible(): + config = _emulator_config() + + settings = sk_learn.SKLearnEmulatorSettings.from_config(config) + assert settings.normalize_y is False + assert settings.random_state is None + assert settings.use_prediction_statistical_uncertainty is False + + config["GPR"]["normalize_y"] = True + config["GPR"]["random_state"] = 20260728 + config["GPR"]["use_prediction_statistical_uncertainty"] = True + settings = sk_learn.SKLearnEmulatorSettings.from_config(config) + assert settings.normalize_y is True + assert settings.random_state == 20260728 + assert settings.use_prediction_statistical_uncertainty is True + + +def test_normalized_gp_restores_target_scale_in_mean_and_covariance(): + settings = sk_learn.SKLearnEmulatorSettings.from_config( + _emulator_config(normalize_y=True) + ) + kernel = Matern( + length_scale=0.4, + length_scale_bounds="fixed", + nu=1.5, + ) + WhiteKernel( + noise_level=0.03, + noise_level_bounds="fixed", + ) + x = np.linspace(0.0, 1.0, 12)[:, np.newaxis] + y = np.sin(2 * np.pi * x[:, 0]) + x_predict = np.array([[0.15], [0.55], [0.9]]) + + base = sk_learn._build_gaussian_process(kernel, settings).fit(x, y) + scaled = sk_learn._build_gaussian_process(kernel, settings).fit(x, 25 * y) + mean, covariance = base.predict(x_predict, return_cov=True) + scaled_mean, scaled_covariance = scaled.predict(x_predict, return_cov=True) + + np.testing.assert_allclose(scaled_mean, 25 * mean) + np.testing.assert_allclose(scaled_covariance, 25**2 * covariance) + + +def test_training_metadata_detects_settings_and_array_changes(tmp_path): + y = np.array([[1.0, 2.0], [2.0, 4.0], [3.0, 6.0]]) + y_err_stat = np.full_like(y, 0.1) + design = np.array([[0.1], [0.5], [0.9]]) + normalized = sk_learn.SKLearnEmulatorSettings.from_config( + _emulator_config(normalize_y=True) + ) + error_aware = sk_learn.SKLearnEmulatorSettings.from_config( + _emulator_config( + normalize_y=True, + use_prediction_statistical_uncertainty=True, + ) + ) + historical = sk_learn.SKLearnEmulatorSettings.from_config( + _emulator_config(normalize_y=False) + ) + analysis_settings = _analysis_settings() + + expected = sk_learn._training_metadata( + normalized, + analysis_settings, + y, + design, + ) + assert expected != sk_learn._training_metadata( + historical, + analysis_settings, + y, + design, + ) + assert expected != sk_learn._training_metadata( + normalized, + analysis_settings, + y + 0.01, + design, + ) + assert expected != sk_learn._training_metadata( + error_aware, + analysis_settings, + y, + design, + y_err_stat, + ) + assert sk_learn._training_metadata( + error_aware, + analysis_settings, + y, + design, + y_err_stat, + ) != sk_learn._training_metadata( + error_aware, + analysis_settings, + y, + design, + y_err_stat + 0.01, + ) + sk_learn._validate_cached_emulator( + pickle.loads(pickle.dumps({"training_metadata": expected})), + expected, + tmp_path / "emulator.pkl", + ) + + with pytest.raises(RuntimeError, match="different settings or input arrays"): + sk_learn._validate_cached_emulator( + {"training_metadata": expected}, + sk_learn._training_metadata( + historical, + analysis_settings, + y, + design, + ), + tmp_path / "emulator.pkl", + ) + + +def test_unverifiable_legacy_cache_is_rejected(tmp_path): + with pytest.raises(RuntimeError, match="no training metadata"): + sk_learn._validate_cached_emulator( + {"PCA": {}, "emulators": []}, + {"version": 1}, + tmp_path / "emulator.pkl", + ) + + +def test_parameterization_changes_invalidate_training_metadata(): + settings = sk_learn.SKLearnEmulatorSettings.from_config(_emulator_config()) + y = np.array([[1.0], [2.0], [3.0]]) + design = np.array([[0.1], [0.5], [0.9]]) + original = _analysis_settings() + changed = _analysis_settings() + changed.raw_analysis_config["parameterization"]["unit"]["max"] = [2.0] + + assert sk_learn._training_metadata( + settings, + original, + y, + design, + ) != sk_learn._training_metadata( + settings, + changed, + y, + design, + ) + + +def test_additional_name_positional_argument_remains_compatible(): + config = _emulator_config() + base_settings = emulation_base.BaseEmulatorSettings.from_emulator_settings( + config + ) + settings = sk_learn.SKLearnEmulatorSettings( + base_settings, + 1, + None, + {"matern": config["kernels"]["matern"]}, + 0, + 1.0e-8, + config, + "legacy_name", + ) + + assert settings.additional_name == "legacy_name" + assert settings.normalize_y is False + + +def test_prediction_uncertainties_are_projected_to_normalized_pc_variance(): + y = np.array( + [ + [1.0, 10.0], + [2.0, 14.0], + [4.0, 18.0], + [8.0, 22.0], + ] + ) + y_err_stat = np.array( + [ + [0.1, 0.4], + [0.2, 0.5], + [0.3, 0.6], + [0.4, 0.7], + ] + ) + scaler = StandardScaler() + pca = PCA(n_components=2, svd_solver="full") + y_pca = pca.fit_transform(scaler.fit_transform(y)) + jitter = 1.0e-8 + + alpha = sk_learn._project_prediction_uncertainties_to_pc_space( + Y_err_stat=y_err_stat, + scaler=scaler, + pca=pca, + Y_pca_truncated=y_pca, + normalize_y=True, + jitter=jitter, + ) + + scaled_variance = (y_err_stat / scaler.scale_) ** 2 + expected = scaled_variance @ (pca.components_**2).T + expected /= np.std(y_pca, axis=0) ** 2 + expected += jitter + np.testing.assert_allclose(alpha, expected) + assert sk_learn._project_prediction_uncertainties_to_pc_space( + Y_err_stat=None, + scaler=scaler, + pca=pca, + Y_pca_truncated=y_pca, + normalize_y=True, + jitter=jitter, + ) is None + + +@pytest.mark.filterwarnings("ignore::sklearn.exceptions.ConvergenceWarning") +def test_fit_emulator_uses_prediction_statistical_uncertainty(tmp_path): + observable = "5020__PbPb__hadron__pt_ch_cms____0-5" + design = np.linspace(0.0, 1.0, 10)[:, np.newaxis] + y = np.vstack( + [ + 0.4 + 0.2 * design[:, 0], + 0.7 - 0.1 * design[:, 0] ** 2, + ] + ) + y_err_stat = np.vstack( + [ + np.linspace(0.01, 0.03, design.shape[0]), + np.linspace(0.02, 0.04, design.shape[0]), + ] + ) + dicttoh5( + { + "Design": design, + "Prediction": { + observable: { + "y": y, + "y_err_stat": y_err_stat, + } + }, + }, + str(tmp_path / "observables.h5"), + ) + analysis_settings = _analysis_settings(tmp_path) + config = _emulator_config( + normalize_y=True, + use_prediction_statistical_uncertainty=True, + ) + config["observable_list"] = [observable] + settings = sk_learn.SKLearnEmulatorSettings.from_config(config) + + result = sk_learn.fit_emulator(settings, analysis_settings) + + alpha_per_pc = result["PCA"]["alpha_per_pc"] + assert alpha_per_pc.shape == (design.shape[0], settings.n_pc) + assert np.ptp(alpha_per_pc[:, 0]) > 0 + np.testing.assert_allclose(result["emulators"][0].alpha, alpha_per_pc[:, 0]) + assert result["training_metadata"]["settings"][ + "use_prediction_statistical_uncertainty" + ] + + emulation_base.IO.write_emulator( + emulator_output=result, + emulator_settings=settings, + analysis_settings=analysis_settings, + ) + assert sk_learn.fit_emulator(settings, analysis_settings) == {} + + analysis_settings.raw_analysis_config["parameterization"]["unit"]["max"] = [ + 2.0 + ] + with pytest.raises(RuntimeError, match="different settings or input arrays"): + sk_learn.fit_emulator(settings, analysis_settings)