autoemulate 2.0.0__tar.gz → 2.1.1__tar.gz

{autoemulate-2.0.0 → autoemulate-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: autoemulate
-Version: 2.0.0
+Version: 2.1.1
 Summary: A python package for semi-automated emulation
 Author: AutoEmulate contributors (see our GitHub page)
 License-Expression: MIT
@@ -39,6 +39,7 @@ Requires-Dist: torchrbf>=0.0.1
 Requires-Dist: arviz>=0.21.0
 Requires-Dist: getdist>=1.7.2
 Requires-Dist: einops>=0.8.1
+Requires-Dist: harmonic>=1.3.0
 Provides-Extra: docs
 Requires-Dist: furo>=2023.9.10; extra == "docs"
 Requires-Dist: sphinx-copybutton>=0.5.2; extra == "docs"
@@ -59,12 +60,13 @@ Requires-Dist: coverage>=7.6.4; extra == "dev"
 Requires-Dist: plotnine>=0.13.6; extra == "dev"
 Requires-Dist: pyright==1.1.405; extra == "dev"
 Requires-Dist: ruff==0.12.11; extra == "dev"
+Requires-Dist: pytest-xdist>=3.0.0; extra == "dev"
 Provides-Extra: spatiotemporal
 Requires-Dist: the-well>=1.1.0; extra == "spatiotemporal"
 Requires-Dist: neuraloperator>=1.0.2; extra == "spatiotemporal"
 Dynamic: license-file
 
-# AutoEmulate <a href="https://alan-turing-institute.github.io/autoemulate/"><img src="https://raw.githubusercontent.com/alan-turing-institute/autoemulate/refs/heads/main/misc/AE_logo_final.png" align="right" height="138" /></a>
+# AutoEmulate <a href="https://alan-turing-institute.github.io/autoemulate/"><img src="https://raw.githubusercontent.com/alan-turing-institute/autoemulate/refs/heads/main/misc/AE.png" align="right" height="138" /></a>
 
 ![CI](https://github.com/alan-turing-institute/autoemulate/actions/workflows/ci.yaml/badge.svg)
 [![codecov](https://codecov.io/gh/alan-turing-institute/autoemulate/graph/badge.svg?token=XD1HXQUIGK)](https://codecov.io/gh/alan-turing-institute/autoemulate)
@@ -139,6 +141,10 @@ You can find the project documentation [here](https://alan-turing-institute.gith
     </tr>
     <tr>
       <td align="center" valign="top" width="14.28%"><a href="https://www.linkedin.com/in/lcm98/"><img src="https://avatars.githubusercontent.com/u/102757707?v=4?s=100" width="100px;" alt="Loh Chun Mun"/><br /><sub><b>Loh Chun Mun</b></sub></a><br /><a href="#doc-cmsamaaa" title="Documentation">📖</a> <a href="#code-cmsamaaa" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/z7076"><img src="https://avatars.githubusercontent.com/u/69681924?v=4?s=100" width="100px;" alt="Jichen Zhang"/><br /><sub><b>Jichen Zhang</b></sub></a><br /><a href="#code-z7076" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/RemDelaporteMathurin"><img src="https://avatars.githubusercontent.com/u/40028739?v=4?s=100" width="100px;" alt="Rémi Delaporte-Mathurin"/><br /><sub><b>Rémi Delaporte-Mathurin</b></sub></a><br /><a href="#bug-RemDelaporteMathurin" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/pmeier"><img src="https://avatars.githubusercontent.com/u/6849766?v=4?s=100" width="100px;" alt="Philip Meier"/><br /><sub><b>Philip Meier</b></sub></a><br /><a href="#bug-pmeier" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/penelopeysm"><img src="https://avatars.githubusercontent.com/u/122629585?v=4?s=100" width="100px;" alt="Penelope Yong"/><br /><sub><b>Penelope Yong</b></sub></a><br /><a href="#ideas-penelopeysm" title="Ideas, Planning, & Feedback">🤔</a> <a href="#review-penelopeysm" title="Reviewed Pull Requests">👀</a> <a href="#code-penelopeysm" title="Code">💻</a> <a href="#doc-penelopeysm" title="Documentation">📖</a></td>
     </tr>
   </tbody>
 </table>

{autoemulate-2.0.0 → autoemulate-2.1.1}/README.md

@@ -1,4 +1,4 @@
-# AutoEmulate <a href="https://alan-turing-institute.github.io/autoemulate/"><img src="https://raw.githubusercontent.com/alan-turing-institute/autoemulate/refs/heads/main/misc/AE_logo_final.png" align="right" height="138" /></a>
+# AutoEmulate <a href="https://alan-turing-institute.github.io/autoemulate/"><img src="https://raw.githubusercontent.com/alan-turing-institute/autoemulate/refs/heads/main/misc/AE.png" align="right" height="138" /></a>
 
 ![CI](https://github.com/alan-turing-institute/autoemulate/actions/workflows/ci.yaml/badge.svg)
 [![codecov](https://codecov.io/gh/alan-turing-institute/autoemulate/graph/badge.svg?token=XD1HXQUIGK)](https://codecov.io/gh/alan-turing-institute/autoemulate)
@@ -73,6 +73,10 @@ You can find the project documentation [here](https://alan-turing-institute.gith
     </tr>
     <tr>
       <td align="center" valign="top" width="14.28%"><a href="https://www.linkedin.com/in/lcm98/"><img src="https://avatars.githubusercontent.com/u/102757707?v=4?s=100" width="100px;" alt="Loh Chun Mun"/><br /><sub><b>Loh Chun Mun</b></sub></a><br /><a href="#doc-cmsamaaa" title="Documentation">📖</a> <a href="#code-cmsamaaa" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/z7076"><img src="https://avatars.githubusercontent.com/u/69681924?v=4?s=100" width="100px;" alt="Jichen Zhang"/><br /><sub><b>Jichen Zhang</b></sub></a><br /><a href="#code-z7076" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/RemDelaporteMathurin"><img src="https://avatars.githubusercontent.com/u/40028739?v=4?s=100" width="100px;" alt="Rémi Delaporte-Mathurin"/><br /><sub><b>Rémi Delaporte-Mathurin</b></sub></a><br /><a href="#bug-RemDelaporteMathurin" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/pmeier"><img src="https://avatars.githubusercontent.com/u/6849766?v=4?s=100" width="100px;" alt="Philip Meier"/><br /><sub><b>Philip Meier</b></sub></a><br /><a href="#bug-pmeier" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/penelopeysm"><img src="https://avatars.githubusercontent.com/u/122629585?v=4?s=100" width="100px;" alt="Penelope Yong"/><br /><sub><b>Penelope Yong</b></sub></a><br /><a href="#ideas-penelopeysm" title="Ideas, Planning, & Feedback">🤔</a> <a href="#review-penelopeysm" title="Reviewed Pull Requests">👀</a> <a href="#code-penelopeysm" title="Code">💻</a> <a href="#doc-penelopeysm" title="Documentation">📖</a></td>
     </tr>
   </tbody>
 </table>

{autoemulate-2.0.0 → autoemulate-2.1.1}/autoemulate/__init__.py

@@ -1,3 +1,6 @@
 from .core.compare import AutoEmulate
+from .core.logging_config import get_logger
+
+get_logger()
 
 __all__ = ["AutoEmulate"]

{autoemulate-2.0.0 → autoemulate-2.1.1}/autoemulate/calibration/base.py

@@ -1,4 +1,3 @@
-import logging
 from collections.abc import Callable
 
 import arviz as az
@@ -7,13 +6,15 @@ from getdist import MCSamples
 from pyro.infer import HMC, MCMC, NUTS, Predictive
 from pyro.infer.mcmc import RandomWalkKernel
 
+from autoemulate.core.logging_config import get_logger
 from autoemulate.core.types import TensorLike
 
+logger = get_logger(__name__)
+
 
 class BayesianMixin:
     """Mixin class for Bayesian calibration methods."""
 
-    logger: logging.Logger
     model: Callable
     observations: dict[str, TensorLike] | None
 
@@ -28,12 +29,12 @@ class BayesianMixin:
         model_kwargs = model_kwargs or {}
         sampler = sampler.lower()
         if sampler == "nuts":
-            self.logger.debug("Using NUTS kernel.")
+            logger.debug("Using NUTS kernel.")
             return NUTS(self.model, **sampler_kwargs)
         if sampler == "hmc":
             step_size = sampler_kwargs.pop("step_size", 0.01)
             trajectory_length = sampler_kwargs.pop("trajectory_length", 1.0)
-            self.logger.debug(
+            logger.debug(
                 "Using HMC kernel with step_size=%s, trajectory_length=%s",
                 step_size,
                 trajectory_length,
@@ -45,9 +46,9 @@ class BayesianMixin:
                 **sampler_kwargs,
             )
         if sampler == "metropolis":
-            self.logger.debug("Using Metropolis (RandomWalkKernel).")
+            logger.debug("Using Metropolis (RandomWalkKernel).")
             return RandomWalkKernel(self.model, **sampler_kwargs)
-        self.logger.error("Unknown sampler: %s", sampler)
+        logger.error("Unknown sampler: %s", sampler)
         raise ValueError(f"Unknown sampler: {sampler}")
 
     def run_mcmc(
@@ -96,11 +97,9 @@ class BayesianMixin:
                         "An initial value must be provided for each chain, parameter "
                         f"{param} tensor only has {init_vals.shape[0]} values."
                     )
-                    self.logger.error(msg)
+                    logger.error(msg)
                     raise ValueError(msg)
-            self.logger.debug(
-                "Initial parameters provided for MCMC: %s", initial_params
-            )
+            logger.debug("Initial parameters provided for MCMC: %s", initial_params)
 
         # Run NUTS
         kernel = self._get_kernel(sampler, model_kwargs=model_kwargs, **sampler_kwargs)
@@ -114,9 +113,9 @@ class BayesianMixin:
             # Multiprocessing
             mp_context="spawn" if num_chains > 1 else None,
         )
-        self.logger.info("Starting MCMC run.")
+        logger.info("Starting MCMC run.")
         mcmc.run()
-        self.logger.info("MCMC run completed.")
+        logger.info("MCMC run completed.")
         return mcmc
 
     def posterior_predictive(self, mcmc: MCMC) -> dict[str, TensorLike]:
@@ -136,7 +135,7 @@ class BayesianMixin:
         posterior_samples = mcmc.get_samples()
         posterior_predictive = Predictive(self.model, posterior_samples)
         samples = posterior_predictive(predict=True)
-        self.logger.debug("Posterior predictive samples generated.")
+        logger.debug("Posterior predictive samples generated.")
         return samples
 
     def to_arviz(
@@ -158,13 +157,13 @@ class BayesianMixin:
         """
         pp_samples = None
         if posterior_predictive:
-            self.logger.info("Including posterior predictive samples in Arviz output.")
+            logger.info("Including posterior predictive samples in Arviz output.")
             pp_samples = self.posterior_predictive(mcmc)
 
         # Need to create dataset manually for Metropolis Hastings
         # This is because az.from_pyro expects kernel with `divergences`
         if isinstance(mcmc.kernel, RandomWalkKernel):
-            self.logger.debug(
+            logger.debug(
                 "Using manual conversion for Metropolis (RandomWalkKernel) kernel."
             )
             if posterior_predictive:
@@ -173,7 +172,7 @@ class BayesianMixin:
                         "Observations must be provided to include observed_data in "
                         "Arviz InferenceData."
                     )
-                    self.logger.error(msg)
+                    logger.error(msg)
                     raise ValueError(msg)
                 az_data = az.InferenceData(
                     posterior=az.convert_to_dataset(
@@ -189,10 +188,10 @@ class BayesianMixin:
                     ),
                 )
         else:
-            self.logger.debug("Using az.from_pyro for conversion.")
+            logger.debug("Using az.from_pyro for conversion.")
             az_data = az.from_pyro(mcmc, posterior_predictive=pp_samples)
 
-        self.logger.info("Arviz InferenceData conversion complete.")
+        logger.info("Arviz InferenceData conversion complete.")
         return az_data
 
     @staticmethod

{autoemulate-2.0.0 → autoemulate-2.1.1}/autoemulate/calibration/bayes.py

@@ -1,13 +1,19 @@
+from collections.abc import Callable
+
 import pyro
 import pyro.distributions as dist
 import torch
+from pyro import poutine
+from pyro.infer import MCMC
 
 from autoemulate.calibration.base import BayesianMixin
 from autoemulate.core.device import TorchDeviceMixin
-from autoemulate.core.logging_config import get_configured_logger
+from autoemulate.core.logging_config import _warn_deprecated_log_level, get_logger
 from autoemulate.core.types import DeviceLike, TensorLike
 from autoemulate.emulators.base import Emulator
 
+logger = get_logger(__name__)
+
 
 class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
     """
@@ -27,7 +33,7 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
         model_discrepancy: float = 0.0,
         calibration_params: list[str] | None = None,
         device: DeviceLike | None = None,
-        log_level: str = "progress_bar",
+        log_level: str | None = None,
     ):
         """
         Initialize the HMC calibration object.
@@ -58,14 +64,8 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
             The device to use. If None, the default torch device is returned.
             TODO: do we need to do anything more to ensure the device is correctly
             handled for the pyro model?
-        log_level: str
-            Logging level for the calibration. Can be one of:
-            - "progress_bar": shows a progress bar during batch simulations
-            - "debug": shows debug messages
-            - "info": shows informational messages
-            - "warning": shows warning messages
-            - "error": shows error messages
-            - "critical": shows critical messages
+        log_level: str | None
+            Deprecated. Configure logging in the calling application instead.
 
         Notes
         -----
@@ -81,10 +81,10 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
             calibration_params = list(parameter_range.keys())
         self.calibration_params = calibration_params
         self.emulator = emulator
-        self.emulator.device = self.device
+        self.emulator.to(self.device)
         self.output_names = list(observations.keys())
-        self.logger, self.progress_bar = get_configured_logger(log_level)
-        self.logger.info(
+        _warn_deprecated_log_level(log_level)
+        logger.info(
             "Initializing BayesianCalibration with parameters: %s",
             self.calibration_params,
         )
@@ -95,12 +95,12 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
         for output, obs in observations.items():
             if obs.ndim == 0:
                 corrected_obs = obs.unsqueeze(0)
-                self.logger.debug(
+                logger.debug(
                     "Observation for output '%s' converted from 0D to 1D.",
                     output,
                 )
             elif obs.ndim > 1:
-                self.logger.error("Tensor for output '%s' is not 1D.", output)
+                logger.error("Tensor for output '%s' is not 1D.", output)
                 raise ValueError(f"Tensor for output '{output}' is not 1D.")
             else:
                 corrected_obs = obs
@@ -108,11 +108,11 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
             obs_lengths.append(corrected_obs.shape[0])
         if len(set(obs_lengths)) != 1:
             msg = "All outputs must have the same number of observations."
-            self.logger.error(msg)
+            logger.error(msg)
             raise ValueError(msg)
         self.observations = processed_observations
         self.n_observations = obs_lengths[0]
-        self.logger.info("Processed observations for outputs: %s", self.output_names)
+        logger.info("Processed observations for outputs: %s", self.output_names)
 
         # Save observation noise as {output: value} dictionary
         if isinstance(observation_noise, float):
@@ -120,12 +120,12 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
                 self.output_names,
                 torch.tensor(observation_noise).to(self.device),
             )
-            self.logger.debug(
+            logger.debug(
                 "Observation noise (variance) set as float: %s", observation_noise
             )
         elif isinstance(observation_noise, dict):
             self.observation_noise = observation_noise
-            self.logger.debug(
+            logger.debug(
                 "Observation noise (variance) set as dict: %s", observation_noise
             )
         else:
@@ -133,7 +133,7 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
                 "Observation noise (variance) must be either a float or a dictionary "
                 "of floats."
             )
-            self.logger.error(msg)
+            logger.error(msg)
             raise ValueError(msg)
 
         self.model_uncertainty = model_uncertainty
@@ -186,8 +186,9 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
             total_variance = (
                 observation_variance + self.model_discrepancy + variance[0, i]
                 if self.model_uncertainty and variance is not None
-                else torch.tensor(observation_variance + self.model_discrepancy).to(
-                    self.device
+                else torch.as_tensor(
+                    observation_variance + self.model_discrepancy,
+                    device=self.device,
                 )
             )
             # Take sqrt for final scale (stddev)
@@ -207,3 +208,137 @@ class BayesianCalibration(TorchDeviceMixin, BayesianMixin):
                         output,
                         dist.Normal(mean[0, i], scale),
                     )
+
+
+def extract_log_probabilities(
+    mcmc: MCMC,
+    model: Callable,
+    device: DeviceLike | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Extract log probabilities from MCMC samples for evidence computation.
+
+    This function extracts posterior samples from a Pyro MCMC object and computes
+    the log probability of each sample under the given probabilistic model. The
+    results are formatted for use with evidence estimation methods like Harmonic.
+
+    Parameters
+    ----------
+    mcmc : MCMC
+        Fitted Pyro MCMC object containing posterior samples. The MCMC object
+        should have been run with multiple chains for best results.
+    model : Callable
+        The Pyro probabilistic model used in MCMC sampling. This should be the
+        same model function passed to the MCMC kernel during sampling.
+    device : DeviceLike | None, optional
+        Device for tensor operations (e.g., 'cpu', 'cuda'). If None, uses the
+        default device. Default is None.
+
+    Returns
+    -------
+    tuple[torch.Tensor, torch.Tensor]
+        A tuple containing:
+        - samples: Tensor of shape (num_chains, num_samples_per_chain, ndim)
+          containing the posterior samples with parameters stacked in the last
+          dimension.
+        - log_probs: Tensor of shape (num_chains, num_samples_per_chain)
+          containing the log probability of each sample under the model.
+
+    Raises
+    ------
+    ValueError
+        If the MCMC object has no samples or if sample extraction fails.
+    RuntimeError
+        If log probability computation fails for any sample.
+
+    Notes
+    -----
+    This function performs the following steps:
+    1. Extracts samples from the MCMC object grouped by chain
+    2. For each sample, conditions the model on the sampled parameter values
+    3. Traces the conditioned model to compute log probabilities
+    4. Returns samples and log probabilities in a format suitable for Harmonic
+
+    The log probabilities include contributions from both the prior and likelihood,
+    representing the unnormalized posterior density at each sample point.
+
+    Examples
+    --------
+    >>> from autoemulate.calibration import (
+    ...     BayesianCalibration,
+    ...     extract_log_probabilities,
+    ... )
+    >>> # After running MCMC calibration
+    >>> bc = BayesianCalibration(emulator, param_range, observations)
+    >>> mcmc = bc.run_mcmc(num_samples=1000, num_chains=4)
+    >>> samples, log_probs = extract_log_probabilities(mcmc, bc.model)
+    >>> print(samples.shape)  # (4, 1000, 2) for 2 parameters
+    >>> print(log_probs.shape)  # (4, 1000)
+
+    See Also
+    --------
+    BayesianCalibration : Class for Bayesian calibration with MCMC
+    EvidenceComputation : Class for computing Bayesian evidence
+
+    """
+    # Set device
+    if device is None:
+        device = torch.device("cpu")
+    elif isinstance(device, str):
+        device = torch.device(device)
+
+    # Extract samples grouped by chain
+    try:
+        samples = mcmc.get_samples(group_by_chain=True)
+    except Exception as e:
+        msg = f"Failed to extract samples from MCMC object: {e}"
+        raise ValueError(msg) from e
+
+    if not samples:
+        msg = "The MCMC object contains no samples"
+        raise ValueError(msg)
+
+    # Get dimensions
+    first_param = next(iter(samples.keys()))
+    num_chains = samples[first_param].shape[0]
+    num_samples_per_chain = samples[first_param].shape[1]
+
+    # Compute log probabilities for each sample
+    log_probs_list = []
+
+    for chain_idx in range(num_chains):
+        chain_log_probs = []
+        for sample_idx in range(num_samples_per_chain):
+            # Extract parameter values for this sample
+            sample_params = {k: v[chain_idx, sample_idx] for k, v in samples.items()}
+
+            # Condition the model on these parameter values
+            conditioned_model = pyro.condition(model, data=sample_params)
+
+            # Trace the model to get log probabilities
+            try:
+                trace = poutine.trace(conditioned_model).get_trace()
+                log_prob = trace.log_prob_sum()
+                # Convert to float, handling both tensor and scalar cases
+                if isinstance(log_prob, torch.Tensor):
+                    log_prob_value = log_prob.item()
+                else:
+                    log_prob_value = float(log_prob)
+                chain_log_probs.append(log_prob_value)
+            except Exception as e:
+                msg = (
+                    f"Failed to compute log probability for chain {chain_idx}, "
+                    f"sample {sample_idx}: {e}"
+                )
+                raise RuntimeError(msg) from e
+
+        log_probs_list.append(chain_log_probs)
+
+    # Convert samples to tensor: (num_chains, num_samples_per_chain, ndim)
+    # Stack parameters in the order they appear in the dictionary
+    samples_tensor = torch.stack([samples[k] for k in samples], dim=-1).to(device)
+
+    # Convert log probabilities to tensor: (num_chains, num_samples_per_chain)
+    log_probs_tensor = torch.tensor(log_probs_list, device=device)
+
+    return samples_tensor, log_probs_tensor