Adds bayesian sampling

.readthedocs.yaml

@@ -17,4 +17,4 @@ formats:
 python:
   install:
     - method: pip
-      path: .[docs]
+      path: .[all]

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#6](https://github.com/pybop-team/PyBOP/issues/6) - Adds bayesian sampling class with NUTS example via Pyro + Torch probabilistic libraries.
 - [#203](https://github.com/pybop-team/PyBOP/pull/203) - Adds support for modern Python packaging via a `pyproject.toml` file and configures the `pytest` test runner and `ruff` linter to use their configurations stored as declarative metadata.
 - [#123](https://github.com/pybop-team/PyBOP/issues/123) - Configures scheduled tests to run against the last three PyPI releases of PyBaMM via dynamic GitHub Actions matrix generation.
 - [#187](https://github.com/pybop-team/PyBOP/issues/187) - Adds M1 Github runner to `test_on_push` workflow, updt. self-hosted supported python versions in scheduled tests.
@@ -30,7 +31,6 @@
 - [#114](https://github.com/pybop-team/PyBOP/issues/114) - Adds a SciPy minimize example and logging for non-Pints optimisers
 - [#116](https://github.com/pybop-team/PyBOP/issues/116) - Adds PSO, SNES, XNES, ADAM, and IPropMin optimisers to PintsOptimisers() class
 - [#38](https://github.com/pybop-team/PyBOP/issues/38) - Restructures the Problem classes ahead of adding a design optimisation example
-- [#38](https://github.com/pybop-team/PyBOP/issues/38) - Updates tests and adds a design optimisation example script `spme_max_energy`
 - [#120](https://github.com/pybop-team/PyBOP/issues/120) - Updates the parameterisation test settings including the number of iterations
 - [#145](https://github.com/pybop-team/PyBOP/issues/145) - Reformats Dataset to contain a dictionary and signal into a list of strings
 

examples/scripts/NUTS_MCMC.py

@@ -0,0 +1,60 @@
+import pybop
+import numpy as np
+import plotly.express as px
+
+# Parameter set and model definition
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+model = pybop.lithium_ion.SPMe(parameter_set=parameter_set)
+
+# Fitting parameters
+parameters = [
+    pybop.Parameter(
+        "Negative electrode active material volume fraction",
+        prior=pybop.Normal(0.72, 0.05),
+        bounds=[0.65, 0.8],
+    ),
+    pybop.Parameter(
+        "Positive electrode active material volume fraction",
+        prior=pybop.Normal(0.62, 0.05),
+        bounds=[0.55, 0.7],
+    ),
+]
+
+# Generate data
+sigma = 0.001
+t_eval = np.arange(0, 1200, 2)
+values = model.predict(t_eval=t_eval)
+corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
+print(
+    parameter_set["Negative electrode active material volume fraction"],
+    parameter_set["Positive electrode active material volume fraction"],
+)
+
+# Dataset definition
+dataset = pybop.Dataset(
+    {
+        "Time [s]": t_eval,
+        "Current function [A]": values["Current [A]"].data,
+        "Voltage [V]": corrupt_values,
+    }
+)
+
+# Generate problem, cost function, and optimisation class
+problem = pybop.FittingProblem(model, parameters, dataset)
+
+# Create the sampler and run
+sampler = pybop.BayesianSampler(problem, "NUTS", transform_space=True)
+samples = sampler.run(
+    num_samples=5, warmup_steps=5, num_chains=1
+)  # Change to 500, 500, 4 for real run
+
+# Plotting
+for param in parameters:
+    param_samples = samples[f"{param.name}"]
+    fig = px.histogram(
+        x=param_samples.numpy(),
+        nbins=60,
+        labels={"x": f"{param.name}"},
+        title=f"Posterior Distribution for {param.name}",
+    )
+    fig.show()

pybop/__init__.py

@@ -62,6 +62,11 @@
 #
 from ._optimisation import Optimisation
 
+#
+# Monte Carlo Sampling class
+#
+from ._inference import BayesianSampler
+
 #
 # Optimiser class
 #
@@ -82,7 +87,7 @@
 #
 from .parameters.parameter import Parameter
 from .parameters.parameter_set import ParameterSet
-from .parameters.priors import Gaussian, Uniform, Exponential
+from .parameters.priors import Normal, Beta, Gaussian, Uniform, Exponential
 
 #
 # Problem class

pybop/_inference.py

@@ -0,0 +1,149 @@
+import torch
+import pyro
+import pyro.distributions as dist
+from pyro.infer import MCMC
+
+
+class BayesianSampler:
+    """
+    Performs Bayesian inference using Markov Chain Monte Carlo (MCMC) sampling with Pyro.
+
+    Parameters
+    ----------
+    problem : object
+        The problem instance containing the model and dataset to be sampled from.
+    kernel : str
+        The name of the Pyro sampling kernel to be used for MCMC.
+    parameter_priors : dict
+        A dictionary of prior distributions for the model parameters.
+    """
+
+    def __init__(self, problem, kernel, transform_space=False):
+        self.problem = problem
+        self.signal = problem.signal
+        self.posterior = None
+        self.parameter_priors = {
+            param.name: (param.prior, param.bounds) for param in problem.parameters
+        }
+        self.kernel = getattr(pyro.infer, kernel)(self.define_sampling_model)
+        self.input_params = torch.zeros(len(self.parameter_priors), dtype=torch.float32)
+        self.transform_space = transform_space
+
+    def define_sampling_model(self):
+        """
+        Defines the model for MCMC sampling, including the prior and likelihood.
+        """
+        # Transform the parameters to the desired interval
+        self.transform_to_parameter_space(self.parameter_priors)
+
+        # Predict the values from the model
+        predicted_values = self.problem.evaluate(x=self.input_params.detach().numpy())
+
+        # Convert your predicted and observed data to PyTorch tensors
+        # breakpoint()
+        voltage_obs = torch.tensor(
+            self.problem._dataset[self.signal[0]].data, dtype=torch.float32
+        )
+        voltage_pred = torch.tensor(predicted_values, dtype=torch.float32)
+        # sigma_tensor = torch.tensor(sigma, dtype=torch.float32)
+
+        # Likelihood of the observations
+        pyro.sample(
+            "obs", dist.Normal(voltage_pred, 0.001).to_event(1), obs=voltage_obs
+        )
+
+    def run(self, num_samples=100, warmup_steps=200, num_chains=1):
+        """
+        Run the MCMC algorithm to sample from the posterior distribution.
+
+        Parameters
+        ----------
+        num_samples : int, optional
+            The number of samples to draw from the posterior. (default: 100)
+        warmup_steps : int, optional
+            The number of warm-up steps before sampling begins. (default: 200)
+        num_chains : int, optional
+            The number of MCMC chains to run. (default: 1)
+
+        Returns
+        -------
+        dict
+            A dictionary of samples from the posterior distribution.
+        """
+
+        self.mcmc = MCMC(
+            self.kernel,
+            num_samples=num_samples,
+            warmup_steps=warmup_steps,
+            num_chains=num_chains,
+        )
+        self.mcmc.run()
+        self.samples = self.mcmc.get_samples()
+
+        if self.transform_space:
+            self.transform_to_parameter_space(
+                self.parameter_priors, samples=self.samples
+            )
+
+        return self.samples
+
+    def transform_to_parameter_space(self, parameter_priors, samples=None):
+        """
+        Transforms unconstrained samples to the parameter space defined by the priors.
+
+        Parameters
+        ----------
+        parameter_priors : dict
+            A dictionary containing the parameter names and their corresponding prior distributions.
+        samples : dict, optional
+            A dictionary containing unconstrained samples to be transformed. If None, new samples are drawn. (default: None)
+        """
+        if samples is None:
+            for i, (param_name, (prior, bounds)) in enumerate(parameter_priors.items()):
+                # Sample from the distribution
+                unconstrained = pyro.sample(f"{param_name}_unconstrained", prior)
+
+                # Transform to the desired interval, and store the value
+                self.input_params[i] = self.scale_to_bounds(unconstrained, bounds)
+        else:
+            for i, (param_name, (prior, bounds)) in enumerate(parameter_priors.items()):
+                # Transform to the desired interval, and store the value
+                self.samples.update(
+                    {
+                        param_name: self.scale_to_bounds(
+                            samples[f"{param_name}_unconstrained"], bounds
+                        )
+                    }
+                )
+
+    @staticmethod
+    def scale_to_bounds(unconstrained, bounds):
+        """
+        Transforms a sample from the [0, 1] interval to the desired [lower_bound, upper_bound] interval.
+
+        Parameters
+        ----------
+        unconstrained : float or torch.Tensor
+            The value in the [0, 1] interval to be transformed.
+        lower_bound : float
+            The lower bound of the target interval.
+        upper_bound : float
+            The upper bound of the target interval.
+
+        Returns
+        -------
+        float or torch.Tensor
+            The transformed value in the [lower_bound, upper_bound] interval.
+        """
+        # Scale the unconstrained value to the new bounds
+        constrained = unconstrained * (bounds[1] - bounds[0]) + bounds[0]
+
+        # Clamp the value to ensure it lies within the specified bounds
+        if hasattr(constrained, "clamp"):
+            # If constrained is a torch.Tensor, use torch's clamp method
+            constrained = constrained.clamp(min=bounds[0], max=bounds[1])
+        else:
+            # If constrained is a float, use the min/max built-in functions
+            constrained = max(min(constrained, bounds[1]), bounds[0])
+
+        return constrained

pybop/_problem.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pybop
 
 
 class BaseProblem:
@@ -55,7 +56,11 @@ def __init__(
         if x0 is None:
             self.x0 = np.zeros(self.n_parameters)
             for i, param in enumerate(self.parameters):
-                self.x0[i] = param.rvs(1)
+                if isinstance(param.prior, (pybop.Normal, pybop.Beta)):
+                    self.x0[i] = param.prior.sample().item()
+                else:
+                    self.x0[i] = param.rvs(1)
+
         elif len(x0) != self.n_parameters:
             raise ValueError("x0 dimensions do not match number of parameters")
 

pybop/parameters/parameter.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pybop
 
 
 class Parameter:
@@ -44,7 +45,7 @@ def __init__(self, name, initial_value=None, prior=None, bounds=None):
         if self.lower_bound >= self.upper_bound:
             raise ValueError("Lower bound must be less than upper bound")
 
-    def rvs(self, n_samples):
+    def rvs(self, n_samples=None):
         """
         Draw random samples from the parameter's prior distribution.
 
@@ -61,7 +62,10 @@ def rvs(self, n_samples):
         array-like
             An array of samples drawn from the prior distribution within the parameter's bounds.
         """
-        samples = self.prior.rvs(n_samples)
+        if isinstance(self.prior, (pybop.Normal, pybop.Beta)):
+            samples = self.prior.sample().item()
+        else:
+            samples = self.prior.rvs(n_samples)
 
         # Constrain samples to be within bounds
         offset = self.margin * (self.upper_bound - self.lower_bound)

pybop/parameters/priors.py

@@ -1,4 +1,5 @@
 import scipy.stats as stats
+from pyro import distributions as dist
 
 
 class Gaussian:
@@ -84,6 +85,14 @@ def __repr__(self):
         return f"{self.name}, mean: {self.mean}, sigma: {self.sigma}"
 
 
+class Normal(dist.Normal):
+    pass
+
+
+class Beta(dist.Beta):
+    pass
+
+
 class Uniform:
     """
     Represents a uniform distribution over a specified interval.

pyproject.toml

@@ -47,7 +47,10 @@ dev = [
     "pytest-xdist",
     "ruff",
     ]
-all = ["pybop[plot]"]
+bayesian = [
+    "pyro-ppl>=1.7"
+    ]
+all = ["pybop[plot,bayesian]"]
 
 [tool.setuptools.packages.find]
 include = ["pybop", "pybop.*"]

readthedocs.yaml

@@ -17,4 +17,4 @@ formats:
 python:
   install:
     - method: pip
-      path: .[docs]
+      path: .[all]

setup.py

@@ -0,0 +1,55 @@
+from distutils.core import setup
+import os
+from setuptools import find_packages
+
+# User-friendly description from README.md
+current_directory = os.path.dirname(os.path.abspath(__file__))
+try:
+    with open(os.path.join(current_directory, "README.md"), encoding="utf-8") as f:
+        long_description = f.read()
+except Exception:
+    long_description = ""
+
+# Defines __version__
+root = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(root, "pybop", "version.py")) as f:
+    exec(f.read())
+
+setup(
+    name="pybop",
+    packages=find_packages("."),
+    version=__version__,  # noqa F821
+    license="BSD-3-Clause",
+    description="Python Battery Optimisation and Parameterisation",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/pybop-team/PyBOP",
+    install_requires=[
+        "pybamm>=23.5",
+        "numpy>=1.16",
+        "scipy>=1.3",
+        "pandas>=1.0",
+        "nlopt>=2.6",
+        "pints>=0.5",
+    ],
+    extras_require={
+        "plot": ["plotly>=5.0"],
+        "bayesian": [
+            "pyro-ppl>=1.7"
+        ],  # Here until it becomes promoted to install_requires
+        "all": ["pybop[plot,bayesian,docs]"],
+        "docs": [
+            "sphinx>=6",
+            "pydata-sphinx-theme",
+            "sphinx-autobuild",
+            "sphinx-autoapi",
+            "sphinx_copybutton",
+            "sphinx_favicon",
+            "sphinx_design",
+            "myst-parser",
+        ],
+    },
+    # https://pypi.org/classifiers/
+    classifiers=[],
+    python_requires=">=3.8,<=3.12",
+)