Fix emulator covariance batching and group cache collisions

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(

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)
 

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)

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)