From d1f47f5f614bdaa23e2ce5b32eb6549b8d94db0f Mon Sep 17 00:00:00 2001 From: botjean Date: Tue, 28 Jul 2026 12:30:24 -0600 Subject: [PATCH] Fix emulator covariance and cache isolation --- src/bayesian/emulation/interface.py | 30 ++++++- src/bayesian/emulation/sk_learn.py | 13 +-- tests/test_emulation_interface.py | 132 ++++++++++++++++++++++++++++ tests/test_emulation_sk_learn.py | 86 ++++++++++++++++++ 4 files changed, 252 insertions(+), 9 deletions(-) create mode 100644 tests/test_emulation_interface.py create mode 100644 tests/test_emulation_sk_learn.py diff --git a/src/bayesian/emulation/interface.py b/src/bayesian/emulation/interface.py index 3bab564..c3d61cd 100644 --- a/src/bayesian/emulation/interface.py +++ b/src/bayesian/emulation/interface.py @@ -280,10 +280,32 @@ def from_config_file( """ c = cls(analysis_settings=analysis_settings) # Initialize the config for each emulator - c.emulation_settings = { - group_name: _emulators[group_cfg["emulator_package"]].EmulatorSettings.from_config(group_cfg) - for group_name, group_cfg in analysis_settings.raw_analysis_config["parameters"]["emulators"].items() - } + c.emulation_settings = {} + group_configs = analysis_settings.raw_analysis_config["parameters"][ + "emulators" + ] + resolved_filenames: dict[str, str] = {} + for group_name, group_cfg in group_configs.items(): + emulator_settings = _emulators[group_cfg["emulator_package"]].EmulatorSettings.from_config(group_cfg) + if "additional_name" in group_cfg: + emulator_settings.additional_name = group_cfg["additional_name"] + elif len(group_configs) > 1: + emulator_settings.additional_name = group_name + + filename = emulation_base.IO.output_filename( + emulator_settings=emulator_settings, + analysis_settings=analysis_settings, + ).name + if filename in resolved_filenames: + other_group = resolved_filenames[filename] + msg = ( + f"Emulator groups '{other_group}' and '{group_name}' both " + f"resolve to '{filename}'. Configure distinct " + "additional_name values." + ) + raise ValueError(msg) + resolved_filenames[filename] = group_name + c.emulation_settings[group_name] = emulator_settings return c def read_all_emulator_groups( diff --git a/src/bayesian/emulation/sk_learn.py b/src/bayesian/emulation/sk_learn.py index 89ca777..886f957 100644 --- a/src/bayesian/emulation/sk_learn.py +++ b/src/bayesian/emulation/sk_learn.py @@ -309,10 +309,11 @@ def predict( emulator_cov_reconstructed_scaled[i_sample] = S.dot(emulator_cov[i_sample].dot(S.T)) assert emulator_cov_reconstructed_scaled.shape == (n_samples, n_features, n_features) - # Include predictive variance due to truncated PCs. - # See comments in mcmc.py for further details. + # Include predictive variance due to truncated PCs. This is a per-parameter-point + # model uncertainty, so it must not depend on how many points are predicted in + # the current batch. for i_sample in range(n_samples): - emulator_cov_reconstructed_scaled[i_sample] += additional_covariance / n_samples + emulator_cov_reconstructed_scaled[i_sample] += additional_covariance # Propagate uncertainty: inverse preprocessing # We only need to undo the unit variance scaling, since the shift does not affect the covariance matrix. @@ -359,8 +360,10 @@ def compute_emulator_cov_unexplained( as a function of theta. We can't do this with the second term, since we didn't emulate it -- so we estimate it, treating it as independent of theta, and add it to the emulator covariance: - Sigma_unexplained = 1/n_samples * S_{>n_pc} D^2_{>n_pc} S_{>n_pc}^T, - where we will include the 1/n_samples factor to account for the fact that we are estimating the covariance from a set of samples. + Sigma_unexplained = S_{>n_pc} V_{>n_pc} S_{>n_pc}^T, + where V is ``PCA.explained_variance_``. Scikit-learn has already normalized + these eigenvalues by ``n_samples - 1``, so no additional division by the + prediction batch size is appropriate. See eqs 21-22 of https://arxiv.org/pdf/2102.11337.pdf TODO: double check this (and compare to https://github.com/jdmulligan/STAT/blob/master/src/emulator.py#L145) diff --git a/tests/test_emulation_interface.py b/tests/test_emulation_interface.py new file mode 100644 index 0000000..f214fd3 --- /dev/null +++ b/tests/test_emulation_interface.py @@ -0,0 +1,132 @@ +from pathlib import Path + +import pytest + +from bayesian import analysis +from bayesian.emulation import base, interface + + +def _config(tmp_path: Path) -> dict: + return { + "observable_table_dir": str(tmp_path), + "observable_config_dir": str(tmp_path), + "observables_filename": "observables.h5", + "output_dir": str(tmp_path / "results"), + "analyses": { + "analysis_test": { + "parameterizations": ["unit"], + "parameterization": { + "unit": { + "names": ["x"], + "min": [0.0], + "max": [1.0], + } + }, + "parameters": { + "emulators": { + "hadron_group": { + "force_retrain": False, + "emulator_package": "sk_learn", + "n_pc": 1, + "kernels": { + "active": ["matern", "noise"], + "matern": { + "nu": 1.5, + "length_scale_bounds_factor": [0.1, 10], + }, + "noise": { + "type": "white", + "args": { + "noise_level": 0.1, + "noise_level_bounds": [0.001, 10], + }, + }, + }, + "GPR": {"n_restarts": 1, "alpha": 1.0e-8}, + "observable_list": ["hadron"], + }, + "jet_group": { + "force_retrain": False, + "emulator_package": "sk_learn", + "n_pc": 1, + "kernels": { + "active": ["matern", "noise"], + "matern": { + "nu": 1.5, + "length_scale_bounds_factor": [0.1, 10], + }, + "noise": { + "type": "white", + "args": { + "noise_level": 0.1, + "noise_level_bounds": [0.001, 10], + }, + }, + }, + "GPR": {"n_restarts": 1, "alpha": 1.0e-8}, + "observable_list": ["jet"], + }, + } + }, + } + }, + } + + +def _analysis_settings(tmp_path: Path, config: dict) -> analysis.AnalysisSettings: + return analysis.AnalysisSettings.from_config( + analysis_name="analysis_test", + config_file=tmp_path / "config.yaml", + config=config, + parameterization="unit", + ) + + +def test_emulator_group_names_produce_distinct_output_filenames(tmp_path: Path) -> None: + config = _config(tmp_path) + analysis_settings = analysis.AnalysisSettings.from_config( + analysis_name="analysis_test", + config_file=tmp_path / "config.yaml", + config=config, + parameterization="unit", + ) + + emulation_config = interface.EmulationConfig.from_config_file(analysis_settings) + filenames = { + group_name: base.IO.output_filename(settings, analysis_settings).name + for group_name, settings in emulation_config.emulation_settings.items() + } + + assert filenames == { + "hadron_group": "emulator_hadron_group.pkl", + "jet_group": "emulator_jet_group.pkl", + } + + +def test_single_unnamed_group_preserves_legacy_filename(tmp_path: Path) -> None: + config = _config(tmp_path) + groups = config["analyses"]["analysis_test"]["parameters"]["emulators"] + groups.pop("jet_group") + analysis_settings = _analysis_settings(tmp_path, config) + + emulation_config = interface.EmulationConfig.from_config_file( + analysis_settings + ) + settings = emulation_config.emulation_settings["hadron_group"] + + assert base.IO.output_filename(settings, analysis_settings).name == "emulator.pkl" + + +@pytest.mark.parametrize("additional_name", ["shared", ""]) +def test_duplicate_explicit_emulator_names_are_rejected( + tmp_path: Path, + additional_name: str, +) -> None: + config = _config(tmp_path) + groups = config["analyses"]["analysis_test"]["parameters"]["emulators"] + for group in groups.values(): + group["additional_name"] = additional_name + analysis_settings = _analysis_settings(tmp_path, config) + + with pytest.raises(ValueError, match="resolve to"): + interface.EmulationConfig.from_config_file(analysis_settings) diff --git a/tests/test_emulation_sk_learn.py b/tests/test_emulation_sk_learn.py new file mode 100644 index 0000000..232815b --- /dev/null +++ b/tests/test_emulation_sk_learn.py @@ -0,0 +1,86 @@ +import numpy as np +from sklearn.decomposition import PCA +from sklearn.preprocessing import StandardScaler + +from bayesian.emulation import sk_learn + + +class _DummySettings: + n_pc = 1 + + +class _ZeroVarianceEmulator: + def predict(self, parameters, return_cov=False): + central = np.zeros(parameters.shape[0]) + if return_cov: + return central, np.zeros((parameters.shape[0], parameters.shape[0])) + return central + + +def test_pca_truncation_covariance_is_batch_invariant(): + y = np.array( + [ + [1.0, 2.0, 3.0], + [2.0, 3.0, 5.0], + [3.0, 5.0, 8.0], + [4.0, 7.0, 11.0], + ] + ) + scaler = StandardScaler() + pca = PCA(n_components=2, svd_solver="full") + pca.fit_transform(scaler.fit_transform(y)) + + additional_covariance = np.array( + [ + [0.10, 0.02, 0.00], + [0.02, 0.20, 0.03], + [0.00, 0.03, 0.30], + ] + ) + results = { + "PCA": { + "pca": pca, + "scaler": scaler, + }, + "emulators": [_ZeroVarianceEmulator()], + } + + parameters = np.array([[0.1, 0.2], [0.3, 0.4]]) + batch_predictions = sk_learn.predict(parameters, results, _DummySettings(), additional_covariance) + single_predictions = [ + sk_learn.predict(parameters[i : i + 1], results, _DummySettings(), additional_covariance) + for i in range(parameters.shape[0]) + ] + + for i, single_prediction in enumerate(single_predictions): + np.testing.assert_allclose(batch_predictions["central_value"][i], single_prediction["central_value"][0]) + np.testing.assert_allclose(batch_predictions["cov"][i], single_prediction["cov"][0]) + + +def test_pca_truncation_covariance_matches_residual_sample_covariance(): + y = np.array( + [ + [1.0, 2.0, 3.0], + [2.0, 3.0, 5.0], + [3.0, 5.0, 8.0], + [4.0, 7.0, 11.0], + [5.0, 11.0, 16.0], + ] + ) + scaler = StandardScaler() + y_scaled = scaler.fit_transform(y) + pca = PCA(svd_solver="full") + y_pca = pca.fit_transform(y_scaled) + retained_scores = np.zeros_like(y_pca) + retained_scores[:, : _DummySettings.n_pc] = y_pca[ + :, : _DummySettings.n_pc + ] + residual = y_scaled - pca.inverse_transform(retained_scores) + expected = np.cov(residual, rowvar=False, ddof=1) + + actual = sk_learn.compute_emulator_cov_unexplained( + _DummySettings(), + {"PCA": {"pca": pca}}, + ) + + np.testing.assert_allclose(actual, expected, atol=1.0e-14)