From d11c8cf280506482a09d97a262cca47aaac598cd Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Fri, 8 Nov 2024 14:33:08 +0000
Subject: [PATCH 1/6] feat: adds annealed importance sampler

---
 examples/scripts/importance_sampling.py |  72 +++++++++++++
 pybop/__init__.py                       |   1 +
 pybop/parameters/priors.py              |  35 +++++++
 pybop/samplers/annealed_importance.py   | 134 ++++++++++++++++++++++++
 4 files changed, 242 insertions(+)
 create mode 100644 examples/scripts/importance_sampling.py
 create mode 100644 pybop/samplers/annealed_importance.py

diff --git a/examples/scripts/importance_sampling.py b/examples/scripts/importance_sampling.py
new file mode 100644
index 000000000..b4639306f
--- /dev/null
+++ b/examples/scripts/importance_sampling.py
@@ -0,0 +1,72 @@
+import numpy as np
+import pybamm
+
+import pybop
+
+# Parameter set and model definition
+solver = pybamm.IDAKLUSolver()
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+parameter_set.update(
+    {
+        "Negative electrode active material volume fraction": 0.66,
+        "Positive electrode active material volume fraction": 0.68,
+    }
+)
+synth_model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solver)
+
+# Fitting parameters
+parameters = pybop.Parameters(
+    pybop.Parameter(
+        "Negative electrode active material volume fraction",
+        prior=pybop.Gaussian(0.65, 0.05),
+    ),
+    pybop.Parameter(
+        "Positive electrode active material volume fraction",
+        prior=pybop.Gaussian(0.65, 0.05),
+    ),
+)
+
+# Generate data
+init_soc = 0.5
+sigma = 0.002
+
+
+def noise(sigma):
+    return np.random.normal(0, sigma, len(values["Voltage [V]"].data))
+
+
+experiment = pybop.Experiment(
+    [
+        (
+            "Discharge at 0.5C for 1 minutes (5 second period)",
+            "Charge at 0.5C for 1 minutes (5 second period)",
+            "Discharge at 3C for 20 seconds (1 second period)",
+        ),
+    ]
+)
+values = synth_model.predict(
+    initial_state={"Initial SoC": init_soc}, experiment=experiment
+)
+
+# Form dataset
+dataset = pybop.Dataset(
+    {
+        "Time [s]": values["Time [s]"].data,
+        "Current function [A]": values["Current [A]"].data,
+        "Voltage [V]": values["Voltage [V]"].data + noise(sigma),
+    }
+)
+
+# Generate problem, likelihood, and sampler
+model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=pybamm.IDAKLUSolver())
+model.build(initial_state={"Initial SoC": init_soc})
+problem = pybop.FittingProblem(model, parameters, dataset)
+likelihood = pybop.GaussianLogLikelihoodKnownSigma(problem, sigma0=sigma)
+prior = pybop.JointLogPrior(*parameters.priors())
+
+sampler = pybop.AnnealedImportanceSampler(
+    likelihood, prior, iterations=10, num_beta=300, cov0=np.eye(2) * 1e-2
+)
+mean, median, std, var = sampler.run()
+
+print(f"mean: {mean}, std: {std}, median: {median}, var: {var}")
diff --git a/pybop/__init__.py b/pybop/__init__.py
index 9eab3d5d6..a280bf900 100644
--- a/pybop/__init__.py
+++ b/pybop/__init__.py
@@ -161,6 +161,7 @@
     SliceRankShrinkingMCMC, SliceStepoutMCMC,
 )
 from .samplers.mcmc_sampler import MCMCSampler
+from .samplers.annealed_importance import AnnealedImportanceSampler
 
 #
 # Observer classes
diff --git a/pybop/parameters/priors.py b/pybop/parameters/priors.py
index 5c40c42c1..b260153cc 100644
--- a/pybop/parameters/priors.py
+++ b/pybop/parameters/priors.py
@@ -430,6 +430,41 @@ def _logpdfS1(self, x: Union[float, np.ndarray]) -> tuple[float, np.ndarray]:
 
         return output, doutput
 
+    def rvs(self, size=1, random_state=None):
+        """
+        Generates random variates from the joint distribution.
+
+        Parameters
+        ----------
+        size : int
+            The number of random variates to generate.
+        random_state : int, optional
+            The random state seed for reproducibility. Default is None.
+
+        Returns
+        -------
+        array_like
+            An array of random variates from the distribution.
+
+        Raises
+        ------
+        ValueError
+            If the size parameter is negative.
+        """
+        if not isinstance(size, (int, tuple)):
+            raise ValueError(
+                "size must be a positive integer or tuple of positive integers"
+            )
+        if isinstance(size, int) and size < 1:
+            raise ValueError("size must be a positive integer")
+        if isinstance(size, tuple) and any(s < 1 for s in size):
+            raise ValueError("size must be a tuple of positive integers")
+
+        samples = []
+        for prior in self._priors:
+            samples.append(prior.rvs(size=size, random_state=random_state)[0])
+        return samples
+
     def __repr__(self) -> str:
         priors_repr = ", ".join([repr(prior) for prior in self._priors])
         return f"{self.__class__.__name__}(priors: [{priors_repr}])"
diff --git a/pybop/samplers/annealed_importance.py b/pybop/samplers/annealed_importance.py
new file mode 100644
index 000000000..c0bdfab2d
--- /dev/null
+++ b/pybop/samplers/annealed_importance.py
@@ -0,0 +1,134 @@
+from typing import Optional
+
+import numpy as np
+
+from pybop import BaseLikelihood, BasePrior
+
+
+class AnnealedImportanceSampler:
+    """
+    This class implements annealed importance sampling of
+    the posterior distribution to compute the model evidence
+    introduced in [1].
+
+    [1] "Annealed Importance Sampling", Radford M. Neal, 1998, Technical Report
+    No. 9805.
+    """
+
+    def __init__(
+        self,
+        log_likelihood: BaseLikelihood,
+        log_prior: BasePrior,
+        x0=None,
+        cov0=None,
+        num_beta: int = 30,
+        iterations: Optional[int] = None,
+    ):
+        self._log_likelihood = log_likelihood
+        self._log_prior = log_prior
+
+        # Set defaults for x0, cov0
+        self._x0 = (
+            x0 if x0 is not None else log_likelihood.parameters.initial_value()
+        )  # Needs transformation
+        self._cov0 = (
+            cov0 if cov0 is not None else np.eye(log_likelihood.n_parameters) * 0.1
+        )
+
+        # Total number of iterations
+        self._iterations = (
+            iterations if iterations is not None else log_likelihood.n_parameters * 1000
+        )
+
+        # Number of beta divisions to consider 0 = beta_n <
+        # beta_n-1 < ... < beta_0 = 1
+        self._num_beta = 30
+
+        self.set_num_beta(num_beta)
+
+    @property
+    def iterations(self) -> int:
+        """Returns the total number of iterations."""
+        return self._iterations
+
+    @iterations.setter
+    def iterations(self, value: int) -> None:
+        """Sets the total number of iterations."""
+        if not isinstance(value, (int, np.integer)):
+            raise TypeError("iterations must be an integer")
+        if value <= 0:
+            raise ValueError("iterations must be positive")
+        self._iterations = int(value)
+
+    @property
+    def num_beta(self) -> int:
+        """Returns the number of beta points"""
+        return self._num_beta
+
+    def set_num_beta(self, num_beta: int) -> None:
+        """Sets the number of beta point values."""
+        if not isinstance(num_beta, (int, np.integer)):
+            raise TypeError("num_beta must be an integer")
+        if num_beta <= 1:
+            raise ValueError("num_beta must be greater than 1")
+        self._num_beta = num_beta
+        self._beta = np.linspace(0, 1, num_beta)
+
+    def run(self) -> tuple[float, float, float, float]:
+        """
+        Run the annealed importance sampling algorithm.
+
+        Returns:
+            Tuple containing (mean, std, median, variance) of the log weights
+
+        Raises:
+            ValueError: If starting position has non-finite log-likelihood
+        """
+        log_w = np.zeros(self._iterations)
+
+        for i in range(self._iterations):
+            current = self._x0.copy()
+
+            if not np.isfinite(self._log_likelihood(current)):
+                raise ValueError("Starting position has non-finite log-likelihood.")
+
+            log_likelihood_current = self._log_likelihood(current)
+            log_prior_current = self._log_prior(current)
+            current_f = log_prior_current + self._beta[0] * log_likelihood_current
+
+            log_density_current = np.zeros(self._num_beta)
+            log_density_previous = np.zeros(self._num_beta)
+            log_density_previous[0] = current_f
+
+            # Main sampling loop
+            for j in range(1, self._num_beta):
+                proposed = np.random.multivariate_normal(current, self._cov0)
+
+                # Evaluate proposed state
+                log_likelihood_proposed = self._log_likelihood(proposed)
+                log_prior_proposed = self._log_prior(proposed)
+
+                # Store proposed
+                log_density_current[j - 1] = current_f
+
+                # Metropolis sampling step
+                if np.isfinite(log_likelihood_proposed):
+                    proposed_f = (
+                        log_prior_proposed + self._beta[j] * log_likelihood_proposed
+                    )
+                    acceptance_log_prob = proposed_f - current_f
+
+                    if np.log(np.random.rand()) < acceptance_log_prob:
+                        current = proposed
+                        current_f = proposed_f
+
+                log_density_previous[j] = current_f
+
+                # Final state
+                log_density_current[self._num_beta - 1] = self._log_prior(
+                    current
+                ) + self._log_likelihood(current)
+            log_w[i] = np.sum(log_density_current) - np.sum(log_density_previous)
+
+        # Return moments of generated chain
+        return np.mean(log_w), np.median(log_w), np.std(log_w), np.var(log_w)

From 5b9974be7e21fa4244473cdef7af91eafc07abb2 Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Sat, 9 Nov 2024 17:24:20 +0000
Subject: [PATCH 2/6] fix: initialise iters from previous val, filter zero vals

---
 pybop/samplers/annealed_importance.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/pybop/samplers/annealed_importance.py b/pybop/samplers/annealed_importance.py
index c0bdfab2d..4fb2fcc67 100644
--- a/pybop/samplers/annealed_importance.py
+++ b/pybop/samplers/annealed_importance.py
@@ -85,10 +85,9 @@ def run(self) -> tuple[float, float, float, float]:
             ValueError: If starting position has non-finite log-likelihood
         """
         log_w = np.zeros(self._iterations)
+        current = self._x0.copy()
 
         for i in range(self._iterations):
-            current = self._x0.copy()
-
             if not np.isfinite(self._log_likelihood(current)):
                 raise ValueError("Starting position has non-finite log-likelihood.")
 
@@ -124,11 +123,12 @@ def run(self) -> tuple[float, float, float, float]:
 
                 log_density_previous[j] = current_f
 
-                # Final state
-                log_density_current[self._num_beta - 1] = self._log_prior(
-                    current
-                ) + self._log_likelihood(current)
+            # Final state
+            log_density_current[self._num_beta - 1] = self._log_prior(
+                current
+            ) + self._log_likelihood(current)
             log_w[i] = np.sum(log_density_current) - np.sum(log_density_previous)
 
-        # Return moments of generated chain
+        # Filter out zeros and return moments of generated chain
+        log_w = log_w[log_w != 0.0]
         return np.mean(log_w), np.median(log_w), np.std(log_w), np.var(log_w)

From 96a26f864eb460fa7ec578171a771488cec1cafa Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Tue, 12 Nov 2024 09:26:54 +0000
Subject: [PATCH 3/6] fix: update algorithm from review

---
 examples/scripts/importance_sampling.py |  9 +++---
 pybop/samplers/annealed_importance.py   | 38 ++++++++++++-------------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/examples/scripts/importance_sampling.py b/examples/scripts/importance_sampling.py
index b4639306f..e1d5f9fcf 100644
--- a/examples/scripts/importance_sampling.py
+++ b/examples/scripts/importance_sampling.py
@@ -18,11 +18,11 @@
 parameters = pybop.Parameters(
     pybop.Parameter(
         "Negative electrode active material volume fraction",
-        prior=pybop.Gaussian(0.65, 0.05),
+        prior=pybop.Gaussian(0.55, 0.02),
     ),
     pybop.Parameter(
         "Positive electrode active material volume fraction",
-        prior=pybop.Gaussian(0.65, 0.05),
+        prior=pybop.Gaussian(0.55, 0.02),
     ),
 )
 
@@ -40,7 +40,6 @@ def noise(sigma):
         (
             "Discharge at 0.5C for 1 minutes (5 second period)",
             "Charge at 0.5C for 1 minutes (5 second period)",
-            "Discharge at 3C for 20 seconds (1 second period)",
         ),
     ]
 )
@@ -58,14 +57,14 @@ def noise(sigma):
 )
 
 # Generate problem, likelihood, and sampler
-model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=pybamm.IDAKLUSolver())
+model = pybop.lithium_ion.SPM(parameter_set=parameter_set, solver=pybamm.IDAKLUSolver())
 model.build(initial_state={"Initial SoC": init_soc})
 problem = pybop.FittingProblem(model, parameters, dataset)
 likelihood = pybop.GaussianLogLikelihoodKnownSigma(problem, sigma0=sigma)
 prior = pybop.JointLogPrior(*parameters.priors())
 
 sampler = pybop.AnnealedImportanceSampler(
-    likelihood, prior, iterations=10, num_beta=300, cov0=np.eye(2) * 1e-2
+    likelihood, prior, chains=10, num_beta=250, cov0=np.eye(2) * 4e-4
 )
 mean, median, std, var = sampler.run()
 
diff --git a/pybop/samplers/annealed_importance.py b/pybop/samplers/annealed_importance.py
index 4fb2fcc67..ed673e9f1 100644
--- a/pybop/samplers/annealed_importance.py
+++ b/pybop/samplers/annealed_importance.py
@@ -22,7 +22,7 @@ def __init__(
         x0=None,
         cov0=None,
         num_beta: int = 30,
-        iterations: Optional[int] = None,
+        chains: Optional[int] = None,
     ):
         self._log_likelihood = log_likelihood
         self._log_prior = log_prior
@@ -36,8 +36,8 @@ def __init__(
         )
 
         # Total number of iterations
-        self._iterations = (
-            iterations if iterations is not None else log_likelihood.n_parameters * 1000
+        self._chains = (
+            chains if chains is not None else log_likelihood.n_parameters * 1000
         )
 
         # Number of beta divisions to consider 0 = beta_n <
@@ -47,18 +47,18 @@ def __init__(
         self.set_num_beta(num_beta)
 
     @property
-    def iterations(self) -> int:
+    def chains(self) -> int:
         """Returns the total number of iterations."""
-        return self._iterations
+        return self._chains
 
-    @iterations.setter
+    @chains.setter
     def iterations(self, value: int) -> None:
         """Sets the total number of iterations."""
         if not isinstance(value, (int, np.integer)):
             raise TypeError("iterations must be an integer")
         if value <= 0:
             raise ValueError("iterations must be positive")
-        self._iterations = int(value)
+        self._chains = int(value)
 
     @property
     def num_beta(self) -> int:
@@ -72,7 +72,7 @@ def set_num_beta(self, num_beta: int) -> None:
         if num_beta <= 1:
             raise ValueError("num_beta must be greater than 1")
         self._num_beta = num_beta
-        self._beta = np.linspace(0, 1, num_beta)
+        self._beta = np.linspace(1, 0, num_beta)
 
     def run(self) -> tuple[float, float, float, float]:
         """
@@ -84,10 +84,10 @@ def run(self) -> tuple[float, float, float, float]:
         Raises:
             ValueError: If starting position has non-finite log-likelihood
         """
-        log_w = np.zeros(self._iterations)
-        current = self._x0.copy()
+        log_w = np.zeros(self._chains)
 
-        for i in range(self._iterations):
+        for i in range(self._chains):
+            current = self._x0.copy()
             if not np.isfinite(self._log_likelihood(current)):
                 raise ValueError("Starting position has non-finite log-likelihood.")
 
@@ -113,22 +113,20 @@ def run(self) -> tuple[float, float, float, float]:
                 # Metropolis sampling step
                 if np.isfinite(log_likelihood_proposed):
                     proposed_f = (
-                        log_prior_proposed + self._beta[j] * log_likelihood_proposed
-                    )
+                        1.0 - self._beta[j]
+                    ) * log_prior_proposed + self._beta[j] * log_likelihood_proposed
                     acceptance_log_prob = proposed_f - current_f
 
                     if np.log(np.random.rand()) < acceptance_log_prob:
                         current = proposed
                         current_f = proposed_f
 
-                log_density_previous[j] = current_f
+                    log_density_previous[j] = (
+                        1.0 - self._beta[j - 1]
+                    ) * log_prior_proposed + self._beta[j - 1] * log_likelihood_proposed
 
-            # Final state
-            log_density_current[self._num_beta - 1] = self._log_prior(
-                current
-            ) + self._log_likelihood(current)
+            # Sum for weights
             log_w[i] = np.sum(log_density_current) - np.sum(log_density_previous)
 
-        # Filter out zeros and return moments of generated chain
-        log_w = log_w[log_w != 0.0]
+        # Return moments of generated chain
         return np.mean(log_w), np.median(log_w), np.std(log_w), np.var(log_w)

From 7bc10001df3f97193b6225ac41c390526925a8e7 Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Wed, 13 Nov 2024 12:45:35 +0000
Subject: [PATCH 4/6] refactor, align equations better with paper

---
 examples/scripts/importance_sampling.py | 11 +++--
 pybop/samplers/annealed_importance.py   | 62 ++++++++++++-------------
 tests/unit/test_annealed_importance.py  | 28 +++++++++++
 tests/unit/test_problem.py              |  2 +-
 4 files changed, 65 insertions(+), 38 deletions(-)
 create mode 100644 tests/unit/test_annealed_importance.py

diff --git a/examples/scripts/importance_sampling.py b/examples/scripts/importance_sampling.py
index e1d5f9fcf..d1c2a1a8d 100644
--- a/examples/scripts/importance_sampling.py
+++ b/examples/scripts/importance_sampling.py
@@ -18,11 +18,11 @@
 parameters = pybop.Parameters(
     pybop.Parameter(
         "Negative electrode active material volume fraction",
-        prior=pybop.Gaussian(0.55, 0.02),
+        prior=pybop.Gaussian(0.6, 0.02),
     ),
     pybop.Parameter(
         "Positive electrode active material volume fraction",
-        prior=pybop.Gaussian(0.55, 0.02),
+        prior=pybop.Gaussian(0.6, 0.02),
     ),
 )
 
@@ -57,14 +57,17 @@ def noise(sigma):
 )
 
 # Generate problem, likelihood, and sampler
-model = pybop.lithium_ion.SPM(parameter_set=parameter_set, solver=pybamm.IDAKLUSolver())
+model = pybop.lithium_ion.SPMe(
+    parameter_set=parameter_set, solver=pybamm.IDAKLUSolver()
+)
 model.build(initial_state={"Initial SoC": init_soc})
 problem = pybop.FittingProblem(model, parameters, dataset)
 likelihood = pybop.GaussianLogLikelihoodKnownSigma(problem, sigma0=sigma)
+# posterior = pybop.LogPosterior(likelihood)
 prior = pybop.JointLogPrior(*parameters.priors())
 
 sampler = pybop.AnnealedImportanceSampler(
-    likelihood, prior, chains=10, num_beta=250, cov0=np.eye(2) * 4e-4
+    likelihood, prior, chains=100, num_beta=100, cov0=np.eye(2) * 4e-4
 )
 mean, median, std, var = sampler.run()
 
diff --git a/pybop/samplers/annealed_importance.py b/pybop/samplers/annealed_importance.py
index ed673e9f1..bcbe0efe9 100644
--- a/pybop/samplers/annealed_importance.py
+++ b/pybop/samplers/annealed_importance.py
@@ -19,31 +19,23 @@ def __init__(
         self,
         log_likelihood: BaseLikelihood,
         log_prior: BasePrior,
-        x0=None,
         cov0=None,
         num_beta: int = 30,
         chains: Optional[int] = None,
     ):
         self._log_likelihood = log_likelihood
         self._log_prior = log_prior
-
-        # Set defaults for x0, cov0
-        self._x0 = (
-            x0 if x0 is not None else log_likelihood.parameters.initial_value()
-        )  # Needs transformation
         self._cov0 = (
             cov0 if cov0 is not None else np.eye(log_likelihood.n_parameters) * 0.1
         )
 
         # Total number of iterations
         self._chains = (
-            chains if chains is not None else log_likelihood.n_parameters * 1000
+            chains if chains is not None else log_likelihood.n_parameters * 300
         )
 
         # Number of beta divisions to consider 0 = beta_n <
         # beta_n-1 < ... < beta_0 = 1
-        self._num_beta = 30
-
         self.set_num_beta(num_beta)
 
     @property
@@ -52,7 +44,7 @@ def chains(self) -> int:
         return self._chains
 
     @chains.setter
-    def iterations(self, value: int) -> None:
+    def chains(self, value: int) -> None:
         """Sets the total number of iterations."""
         if not isinstance(value, (int, np.integer)):
             raise TypeError("iterations must be an integer")
@@ -72,7 +64,15 @@ def set_num_beta(self, num_beta: int) -> None:
         if num_beta <= 1:
             raise ValueError("num_beta must be greater than 1")
         self._num_beta = num_beta
-        self._beta = np.linspace(1, 0, num_beta)
+        self._beta = np.linspace(0, 1, num_beta)
+
+    def transition_distribution(self, x, j):
+        """
+        Annealed Importance Sampling - Equation 3.
+        """
+        return (1.0 - self._beta[j]) * self._log_prior(x) + self._beta[
+            j
+        ] * self._log_likelihood(x)
 
     def run(self) -> tuple[float, float, float, float]:
         """
@@ -87,46 +87,42 @@ def run(self) -> tuple[float, float, float, float]:
         log_w = np.zeros(self._chains)
 
         for i in range(self._chains):
-            current = self._x0.copy()
+            current = self._log_prior.rvs()
             if not np.isfinite(self._log_likelihood(current)):
                 raise ValueError("Starting position has non-finite log-likelihood.")
 
-            log_likelihood_current = self._log_likelihood(current)
-            log_prior_current = self._log_prior(current)
-            current_f = log_prior_current + self._beta[0] * log_likelihood_current
+            current_f = self._log_prior(current)
 
             log_density_current = np.zeros(self._num_beta)
+            log_density_current[0] = current_f
             log_density_previous = np.zeros(self._num_beta)
             log_density_previous[0] = current_f
 
             # Main sampling loop
             for j in range(1, self._num_beta):
-                proposed = np.random.multivariate_normal(current, self._cov0)
+                # Compute transition with current sample
+                log_density_current[j] = self.transition_distribution(current, j)
 
-                # Evaluate proposed state
-                log_likelihood_proposed = self._log_likelihood(proposed)
-                log_prior_proposed = self._log_prior(proposed)
+                # Calculate the previous transition with current sample
+                log_density_previous[j] = self.transition_distribution(current, j - 1)
 
-                # Store proposed
-                log_density_current[j - 1] = current_f
+                # Generate new sample from current (eqn.4)
+                proposed = np.random.multivariate_normal(current, self._cov0)
 
-                # Metropolis sampling step
-                if np.isfinite(log_likelihood_proposed):
-                    proposed_f = (
-                        1.0 - self._beta[j]
-                    ) * log_prior_proposed + self._beta[j] * log_likelihood_proposed
-                    acceptance_log_prob = proposed_f - current_f
+                # Evaluate proposed sample
+                if np.isfinite(self._log_likelihood(proposed)):
+                    proposed_f = self.transition_distribution(proposed, j)
 
+                    # Metropolis acceptance
+                    acceptance_log_prob = proposed_f - current_f
                     if np.log(np.random.rand()) < acceptance_log_prob:
                         current = proposed
                         current_f = proposed_f
 
-                    log_density_previous[j] = (
-                        1.0 - self._beta[j - 1]
-                    ) * log_prior_proposed + self._beta[j - 1] * log_likelihood_proposed
-
-            # Sum for weights
-            log_w[i] = np.sum(log_density_current) - np.sum(log_density_previous)
+            # Sum for weights (eqn.24)
+            log_w[i] = (
+                np.sum(log_density_current - log_density_previous) / self._num_beta
+            )
 
         # Return moments of generated chain
         return np.mean(log_w), np.median(log_w), np.std(log_w), np.var(log_w)
diff --git a/tests/unit/test_annealed_importance.py b/tests/unit/test_annealed_importance.py
new file mode 100644
index 000000000..43954dd0e
--- /dev/null
+++ b/tests/unit/test_annealed_importance.py
@@ -0,0 +1,28 @@
+import numpy as np
+import pytest
+
+import pybop
+
+
+class TestPintsSamplers:
+    """
+    Class for unit tests of AnnealedImportanceSampler.
+    """
+
+    @pytest.mark.unit
+    def test_annealed_importance_sampler(self):
+        likelihood = pybop.Gaussian(5, 0.5)
+
+        def scaled_likelihood(x):
+            return likelihood(x) * 2
+
+        prior = pybop.Gaussian(4.7, 2)
+
+        # Sample
+        sampler = pybop.AnnealedImportanceSampler(
+            scaled_likelihood, prior, chains=15, num_beta=200, cov0=np.eye(1) * 5e-2
+        )
+        mean, median, std, var = sampler.run()
+
+        # Assertions to be added
+        print(f"mean: {mean}, std: {std}, median: {median}, var: {var}")
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index c02a0aa17..12047217a 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -63,7 +63,7 @@ def dataset(self, model, experiment):
 
     @pytest.fixture
     def signal(self):
-        return "Voltage [V]"
+        return ["Voltage [V]"]
 
     @pytest.mark.unit
     def test_base_problem(self, parameters, model, dataset):

From 4f4559110d92ff7a1be2683c9cae0589a20a19c6 Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Wed, 13 Nov 2024 17:52:05 +0000
Subject: [PATCH 5/6] bugfixes for problem tests, updt docstrings

---
 pybop/samplers/annealed_importance.py | 8 ++++----
 tests/unit/test_problem.py            | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/pybop/samplers/annealed_importance.py b/pybop/samplers/annealed_importance.py
index bcbe0efe9..3f0f301eb 100644
--- a/pybop/samplers/annealed_importance.py
+++ b/pybop/samplers/annealed_importance.py
@@ -68,7 +68,7 @@ def set_num_beta(self, num_beta: int) -> None:
 
     def transition_distribution(self, x, j):
         """
-        Annealed Importance Sampling - Equation 3.
+        Transition distribution for each beta value [j] - Eqn 3.
         """
         return (1.0 - self._beta[j]) * self._log_prior(x) + self._beta[
             j
@@ -79,7 +79,7 @@ def run(self) -> tuple[float, float, float, float]:
         Run the annealed importance sampling algorithm.
 
         Returns:
-            Tuple containing (mean, std, median, variance) of the log weights
+            Tuple containing (mean, median, std, variance) of the log weights
 
         Raises:
             ValueError: If starting position has non-finite log-likelihood
@@ -100,7 +100,7 @@ def run(self) -> tuple[float, float, float, float]:
 
             # Main sampling loop
             for j in range(1, self._num_beta):
-                # Compute transition with current sample
+                # Compute jth transition with current sample
                 log_density_current[j] = self.transition_distribution(current, j)
 
                 # Calculate the previous transition with current sample
@@ -124,5 +124,5 @@ def run(self) -> tuple[float, float, float, float]:
                 np.sum(log_density_current - log_density_previous) / self._num_beta
             )
 
-        # Return moments of generated chain
+        # Return moments across chains
         return np.mean(log_w), np.median(log_w), np.std(log_w), np.var(log_w)
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index 12047217a..ff0ecc425 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -248,8 +248,8 @@ def test_multi_fitting_problem(self, model, parameters, dataset, signal):
             problem_1._dataset["Time [s]"]
         ) + len(problem_2._dataset["Time [s]"])
         assert len(combined_problem._dataset["Combined signal"]) == len(
-            problem_1._dataset[signal]
-        ) + len(problem_2._dataset[signal])
+            problem_1._dataset[signal[0]]
+        ) + len(problem_2._dataset[signal[0]])
 
         y = combined_problem.evaluate(inputs=[1e-5, 1e-5])
         assert len(y["Combined signal"]) == len(

From c5eef05cf5ed0db7ba4ad9ae20f03f41a93c5269 Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Thu, 14 Nov 2024 11:20:18 +0000
Subject: [PATCH 6/6] Add integral calculation

---
 pybop/samplers/annealed_importance.py  | 16 +++++++++++++---
 tests/unit/test_annealed_importance.py |  6 +++---
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/pybop/samplers/annealed_importance.py b/pybop/samplers/annealed_importance.py
index 3f0f301eb..026412ab7 100644
--- a/pybop/samplers/annealed_importance.py
+++ b/pybop/samplers/annealed_importance.py
@@ -74,7 +74,7 @@ def transition_distribution(self, x, j):
             j
         ] * self._log_likelihood(x)
 
-    def run(self) -> tuple[float, float, float, float]:
+    def run(self) -> tuple[float, float, float]:
         """
         Run the annealed importance sampling algorithm.
 
@@ -85,6 +85,8 @@ def run(self) -> tuple[float, float, float, float]:
             ValueError: If starting position has non-finite log-likelihood
         """
         log_w = np.zeros(self._chains)
+        I = np.zeros(self._chains)
+        samples = np.zeros(self._num_beta)
 
         for i in range(self._chains):
             current = self._log_prior.rvs()
@@ -119,10 +121,18 @@ def run(self) -> tuple[float, float, float, float]:
                         current = proposed
                         current_f = proposed_f
 
+                samples[j] = current
+
             # Sum for weights (eqn.24)
             log_w[i] = (
                 np.sum(log_density_current - log_density_previous) / self._num_beta
             )
 
-        # Return moments across chains
-        return np.mean(log_w), np.median(log_w), np.std(log_w), np.var(log_w)
+            # Compute integral using weights and samples
+            I[i] = np.mean(
+                self._log_likelihood(samples)
+                * np.exp((log_density_current - log_density_previous) / self._num_beta)
+            )
+
+        # Return log weights, integral, samples
+        return log_w, I, samples
diff --git a/tests/unit/test_annealed_importance.py b/tests/unit/test_annealed_importance.py
index 43954dd0e..08287fb5b 100644
--- a/tests/unit/test_annealed_importance.py
+++ b/tests/unit/test_annealed_importance.py
@@ -20,9 +20,9 @@ def scaled_likelihood(x):
 
         # Sample
         sampler = pybop.AnnealedImportanceSampler(
-            scaled_likelihood, prior, chains=15, num_beta=200, cov0=np.eye(1) * 5e-2
+            scaled_likelihood, prior, chains=15, num_beta=500, cov0=np.eye(1) * 1e-2
         )
-        mean, median, std, var = sampler.run()
+        log_w, I, samples = sampler.run()
 
         # Assertions to be added
-        print(f"mean: {mean}, std: {std}, median: {median}, var: {var}")
+        print(f"Integral: {np.mean(I)}, std: {np.std(I)}")