Base Bayesian framework with EP-BOLFI integration

CHANGELOG.md

@@ -13,6 +13,7 @@
 ## Features
 
 - [#815](https://github.com/pybop-team/PyBOP/pull/815) - Adds function import_pyprobe_result to import a pyprobe.result into a pybop.dataset. Allows for creating a dataset directly from a pybamm.solution object.
+- [#846](https://github.com/pybop-team/PyBOP/pull/846) - Adds Bayesian optimisation framework and, as an example, the EP-BOLFI optimiser
 
 ## Optimisations
 

examples/scripts/battery_parameterisation/bayes_gitt_fitting.py

@@ -0,0 +1,125 @@
+import numpy as np
+import pybamm
+from ep_bolfi.models.solversetup import spectral_mesh_pts_and_method
+from pybamm import CasadiSolver, Experiment, print_citations
+
+import pybop
+from pybop.costs.feature_distances import SquareRootFeatureDistance
+from pybop.optimisers.ep_bolfi_optimiser import EP_BOLFI, EPBOLFIOptions
+from pybop.parameters.multivariate_distributions import MultivariateGaussian
+from pybop.parameters.multivariate_parameters import MultivariateParameters
+
+parameter_set = pybamm.ParameterValues("Chen2020")
+original_D_n = parameter_set["Negative particle diffusivity [m2.s-1]"]
+original_D_p = parameter_set["Positive particle diffusivity [m2.s-1]"]
+
+unknowns = MultivariateParameters(
+    {
+        "Negative particle diffusivity [m2.s-1]": pybop.Parameter(
+            initial_value=0.9 * original_D_n,
+            bounds=[original_D_n / 2, original_D_n * 2],
+            transformation=pybop.LogTransformation(),
+        ),
+        "Positive particle diffusivity [m2.s-1]": pybop.Parameter(
+            initial_value=1.1 * original_D_p,
+            bounds=[original_D_p / 2, original_D_p * 2],
+            transformation=pybop.LogTransformation(),
+        ),
+    },
+    distribution=MultivariateGaussian(
+        [np.log(original_D_n), np.log(original_D_p)],
+        [[np.log(2), 0.0], [0.0, np.log(2)]],
+    ),
+)
+
+submesh_types, var_pts, spatial_methods = spectral_mesh_pts_and_method(10, 10, 10)
+parameter_values = pybamm.ParameterValues("Chen2020")
+# Put empty Parameter slots as placeholders.
+parameter_values["Negative particle diffusivity [m2.s-1]"] = pybop.Parameter()
+parameter_values["Positive particle diffusivity [m2.s-1]"] = pybop.Parameter()
+
+simulator = pybop.pybamm.Simulator(
+    model=pybamm.lithium_ion.DFN(),
+    parameter_values=parameter_values,
+    protocol=Experiment(
+        [
+            "Discharge at 1.0 C for 15 minutes (1 second period)",
+            "Rest for 15 minutes (1 second period)",
+        ]
+    ),
+    solver=CasadiSolver(
+        rtol=1e-5,
+        atol=1e-5,
+        root_tol=1e-3,
+        max_step_decrease_count=10,
+        extra_options_setup={
+            "disable_internal_warnings": True,
+            "newton_scheme": "tfqmr",
+        },
+        return_solution_if_failed_early=True,
+    ),
+    output_variables=["Voltage [V]"],
+    submesh_types=submesh_types,
+    var_pts=var_pts,
+    spatial_methods=spatial_methods,
+)
+# Override the forced univariate Parameters.
+simulator.parameters = unknowns
+
+synthetic_data = simulator.solve(
+    {
+        "Negative particle diffusivity [m2.s-1]": original_D_n,
+        "Positive particle diffusivity [m2.s-1]": original_D_p,
+    }
+)
+
+dataset = pybop.Dataset(
+    {
+        "Time [s]": synthetic_data["Time [s]"].data,
+        "Current function [A]": synthetic_data["Current [A]"].data,
+        "Voltage [V]": synthetic_data["Voltage [V]"].data,
+    }
+)
+
+ICI_cost = SquareRootFeatureDistance(
+    dataset["Time [s]"],
+    dataset["Voltage [V]"],
+    feature="inverse_slope",
+    time_start=0,
+    time_end=90,
+)
+GITT_cost = SquareRootFeatureDistance(
+    dataset["Time [s]"],
+    dataset["Voltage [V]"],
+    feature="inverse_slope",
+    time_start=901,
+    time_end=991,
+)
+
+if __name__ == "__main__":
+    ICI_problem = pybop.Problem(simulator, ICI_cost)
+    GITT_problem = pybop.Problem(simulator, GITT_cost)
+    problem = pybop.MetaProblem(ICI_problem, GITT_problem)
+    # Overwrite the forced pybop.Parameters with MultivariateParameters.
+    problem.parameters = unknowns
+    options = EPBOLFIOptions(
+        ep_iterations=2,
+        ep_total_dampening=0,
+        bolfi_initial_sobol_samples=10,
+        bolfi_optimally_acquired_samples=10,
+        bolfi_posterior_effective_sample_size=20,
+        posterior_gelman_rubin_threshold=1.2,
+        verbose=True,
+    )
+    optim = EP_BOLFI(problem, options)
+
+    result = optim.run()
+
+    import plotly.io as pio
+
+    pio.renderers.default = "browser"
+
+    pybop.plot.convergence(result, yaxis={"type": "log"})
+    pybop.plot.parameters(result, yaxis={"type": "log"}, yaxis2={"type": "log"})
+
+    print_citations()

pybop/__init__.py

@@ -48,7 +48,9 @@
 # Parameter classes
 #
 from .parameters.parameter import Parameter, Parameters
+from .parameters.multivariate_parameters import MultivariateParameters
 from .parameters.priors import BasePrior, Gaussian, Uniform, Exponential, JointPrior
+from .parameters.multivariate_distributions import MultivariateNonparametric, MultivariateUniform, MultivariateGaussian
 
 #
 # Model classes

pybop/_result.py

@@ -5,6 +5,7 @@
 if TYPE_CHECKING:
     from pybop import BaseOptimiser, BaseSampler
 from pybop import Logger, plot
+from pybop.parameters.multivariate_parameters import MultivariateParameters
 
 
 class OptimisationResult:
@@ -117,7 +118,7 @@ def combine(results: list["OptimisationResult"]) -> "OptimisationResult":
 
         ret._best_run = None  # noqa: SLF001
         ret.n_runs = len(results)
-        ret._validate()  #  noqa: SLF001
+        ret._validate()  # noqa: SLF001
 
         return ret
 
@@ -360,3 +361,74 @@ def __init__(
             message=message,
         )
         self.chains = chains
+
+
+class BayesianOptimisationResult(OptimisationResult):
+    """
+    Stores the result of a Bayesian optimisation or a Bayesian model
+    selection.
+
+    Attributes
+    ----------
+    problem: pybop.Problem
+        The optimisation object used to generate the results.
+    x : ndarray
+        The solution of the optimisation (in model space).
+    final_cost : float
+        The cost associated with the solution x.
+    n_iterations : int or dict
+        Number of iterations performed by the optimiser. Since Bayesian
+        optimisers tend to have layers of various optimisation
+        algorithms, their iteration counts may be put individually.
+    n_evaluations : int or dict
+        Number of evaluations performed by the optimiser. Since Bayesian
+        optimisers tend to have layers of various optimisation
+        algorithms, their evaluation counts my be put individually.
+    message : str
+        The reason for stopping given by the optimiser.
+    lower_bounds: ndarray
+        The lower confidence parameter boundaries.
+    upper_bounds: ndarray
+        The upper confidence parameter boundaries.
+    posterior : MultivariateParameters
+        The probability distribution of the optimisation.
+    maximum_a_posteriori : Inputs or ndarray
+        Complementing the best observed value in `x`, this is the
+        prediction for the best parameter value.
+    log_evidence_mean : float
+        The logarithm of the evidence of the parameterization. Higher
+        values are better. May only be interpreted relative to a
+        calibration case, e.g., a test-run with synthetic data.
+    log_evidence_variance : float
+        The logarithm of the variance in the calculation of the
+        evidence. For reliable comparisons based on the evidence, should
+        be at or below the scale of the evidence itself.
+    """
+
+    def __init__(
+        self,
+        optim: "BaseOptimiser",
+        logger: Logger,
+        time: float | dict,
+        optim_name: str | None = None,
+        message: str | None = None,
+        lower_bounds: np.ndarray | None = None,
+        upper_bounds: np.ndarray | None = None,
+        posterior: MultivariateParameters | None = None,
+        maximum_a_posteriori: np.ndarray | None = None,
+        log_evidence_mean: float | None = None,
+        log_evidence_variance: float | None = None,
+    ):
+        super().__init__(
+            optim=optim,
+            logger=logger,
+            time=time,
+            optim_name=optim_name,
+            message=message,
+        )
+        self.lower_bounds = lower_bounds
+        self.upper_bounds = upper_bounds
+        self.posterior = posterior
+        self.maximum_a_posteriori = maximum_a_posteriori
+        self.log_evidence_mean = log_evidence_mean
+        self.log_evidence_variance = log_evidence_variance