Adds JAX IDAKLU solver integration

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#481](https://github.com/pybop-team/PyBOP/pull/481) - Adds experimental support for PyBaMM's jaxified IDAKLU solver. Includes Jax-specific cost functions `pybop.JaxSumSquareError` and `pybop.JaxLogNormalLikelihood`. Adds `Jax` optional dependency to PyBaMM dependency.
 - [#452](https://github.com/pybop-team/PyBOP/issues/452) - Extends `cell_mass` and `approximate_capacity` for half-cell models.
 - [#544](https://github.com/pybop-team/PyBOP/issues/544) - Allows iterative plotting using `StandardPlot`.
 - [#541](https://github.com/pybop-team/PyBOP/pull/541) - Adds `ScaledLogLikelihood` and `BaseMetaLikelihood` classes.

examples/scripts/jax-solver-example.py

@@ -0,0 +1,66 @@
+import numpy as np
+import pybamm
+
+import pybop
+
+# Parameter set and model definition
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+
+# The IDAKLU, and it's jaxified version perform very well on the DFN with and without
+# gradient calculations
+solver = pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6)
+model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solver)
+
+# Fitting parameters
+parameters = pybop.Parameters(
+    pybop.Parameter(
+        "Negative electrode active material volume fraction",
+        initial_value=0.55,
+        bounds=[0.5, 0.8],
+    ),
+    pybop.Parameter(
+        "Positive electrode active material volume fraction",
+        initial_value=0.55,
+        bounds=[0.5, 0.8],
+    ),
+)
+
+# Define test protocol and generate data
+t_eval = np.linspace(0, 600, 600)
+values = model.predict(
+    initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval
+)
+
+# Form dataset
+dataset = pybop.Dataset(
+    {
+        "Time [s]": values["Time [s]"].data,
+        "Current function [A]": values["Current [A]"].data,
+        "Voltage [V]": values["Voltage [V]"].data,
+    }
+)
+
+# Construct the Problem
+problem = pybop.FittingProblem(model, parameters, dataset)
+
+# By selecting a Jax based cost function, the IDAKLU solver will be
+# jaxified (wrapped in a Jax compiled expression) and used for optimisation
+cost = pybop.JaxLogNormalLikelihood(problem, sigma0=2e-3)
+
+# Non-gradient optimiser, change to `pybop.AdamW` for gradient-based example
+optim = pybop.IRPropMin(
+    cost,
+    max_unchanged_iterations=20,
+    max_iterations=100,
+)
+
+results = optim.run()
+
+# Plot convergence
+pybop.plot.convergence(optim)
+
+# Plot parameter trace
+pybop.plot.parameters(optim)
+
+# Plot voronoi optimiser surface
+pybop.plot.surface(optim)

examples/scripts/jaxified-idaklu-benchmarks.py

@@ -0,0 +1,106 @@
+import time
+
+import numpy as np
+import pybamm
+
+import pybop
+
+n = 1  # Number of solves
+solvers = [
+    pybamm.CasadiSolver(mode="fast with events", atol=1e-6, rtol=1e-6),
+    pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6),
+]
+
+# Parameter set and model definition
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solvers[0])
+
+# Fitting parameters
+parameters = pybop.Parameters(
+    pybop.Parameter(
+        "Negative electrode active material volume fraction", initial_value=0.55
+    ),
+    pybop.Parameter(
+        "Positive electrode active material volume fraction", initial_value=0.55
+    ),
+)
+
+# Define test protocol and generate data
+t_eval = np.linspace(0, 100, 1000)
+values = model.predict(
+    initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval
+)
+
+# Form dataset
+dataset = pybop.Dataset(
+    {
+        "Time [s]": values["Time [s]"].data,
+        "Current function [A]": values["Current [A]"].data,
+        "Voltage [V]": values["Voltage [V]"].data,
+    }
+)
+
+
+# Create inputs function for benchmarking
+def inputs():
+    return {
+        "Negative electrode active material volume fraction": 0.55
+        + np.random.normal(0, 0.01),
+        "Positive electrode active material volume fraction": 0.55
+        + np.random.normal(0, 0.01),
+    }
+
+
+# Iterate over the solvers and print benchmarks
+for solver in solvers:
+    # Setup Fitting Problem
+    model.solver = solver
+    problem = pybop.FittingProblem(model, parameters, dataset)
+    cost = pybop.SumSquaredError(problem)
+
+    start_time = time.time()
+    for _i in range(n):
+        out = problem.model.simulate(inputs=inputs(), t_eval=t_eval)
+    print(f"({solver.name}) Time model.simulate: {time.time() - start_time:.4f}")
+
+    start_time = time.time()
+    for _i in range(n):
+        out = problem.model.simulateS1(inputs=inputs(), t_eval=t_eval)
+    print(f"({solver.name}) Time model.SimulateS1: {time.time() - start_time:.4f}")
+
+    start_time = time.time()
+    for _i in range(n):
+        out = problem.evaluate(inputs=inputs())
+    print(f"({solver.name}) Time problem.evaluate: {time.time() - start_time:.4f}")
+
+    start_time = time.time()
+    for _i in range(n):
+        out = problem.evaluateS1(inputs=inputs())
+    print(f"({solver.name}) Time Problem.EvaluateS1: {time.time() - start_time:.4f}")
+
+    start_time = time.time()
+    for _i in range(n):
+        out = cost(inputs(), calculate_grad=False)
+    print(f"({solver.name}) Time PyBOP Cost w/o grad: {time.time() - start_time:.4f}")
+
+    start_time = time.time()
+    for _i in range(n):
+        out = cost(inputs(), calculate_grad=True)
+    print(f"({solver.name}) Time PyBOP Cost w/grad: {time.time() - start_time:.4f}")
+
+# Recreate for Jax IDAKLU solver
+ida_solver = pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6)
+model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=ida_solver)
+problem = pybop.FittingProblem(model, parameters, dataset)
+cost = pybop.JaxSumSquaredError(problem)
+
+# Jaxified benchmarks
+start_time = time.time()
+for _i in range(n):
+    out = cost(inputs(), calculate_grad=False)
+print(f"Time Jax SumSquaredError w/o grad: {time.time() - start_time:.4f}")
+
+start_time = time.time()
+for _i in range(n):
+    out = cost(inputs(), calculate_grad=True)
+print(f"Time Jax SumSquaredError w/ grad: {time.time() - start_time:.4f}")

examples/scripts/maximum_likelihood.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pybamm
 
 import pybop
 
@@ -7,10 +8,12 @@
 parameter_set.update(
     {
         "Negative electrode active material volume fraction": 0.63,
-        "Positive electrode active material volume fraction": 0.51,
+        "Positive electrode active material volume fraction": 0.62,
     }
 )
-model = pybop.lithium_ion.SPM(parameter_set=parameter_set)
+options = {"max_num_steps": int(1e6), "max_error_test_failures": 60}
+solver = pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, options=options)
+model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solver)
 
 # Fitting parameters
 parameters = pybop.Parameters(
@@ -57,8 +60,8 @@ def noise(sigma):
 signal = ["Voltage [V]", "Bulk open-circuit voltage [V]"]
 # Generate problem, cost function, and optimisation class
 problem = pybop.FittingProblem(model, parameters, dataset, signal=signal)
-likelihood = pybop.GaussianLogLikelihood(problem, sigma0=sigma * 4)
-optim = pybop.IRPropMin(
+likelihood = pybop.JaxGaussianLogLikelihoodKnownSigma(problem, sigma0=sigma)
+optim = pybop.XNES(
     likelihood,
     max_unchanged_iterations=20,
     min_iterations=20,

pybop/__init__.py

@@ -117,6 +117,11 @@
 )
 from .costs._weighted_cost import WeightedCost
 
+#
+# Experimental
+#
+from .experimental import BaseJaxCost, JaxSumSquaredError, JaxLogNormalLikelihood, JaxGaussianLogLikelihoodKnownSigma
+
 #
 # Optimiser classes
 #

pybop/experimental/__init__.py

@@ -0,0 +1 @@
+from .jax_costs import BaseJaxCost, JaxLogNormalLikelihood, JaxSumSquaredError, JaxGaussianLogLikelihoodKnownSigma

pybop/experimental/jax_costs.py

@@ -0,0 +1,166 @@
+from typing import Union
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from pybamm import IDAKLUSolver
+
+from pybop import BaseCost, BaseLikelihood, BaseProblem, Inputs
+
+
+class BaseJaxCost(BaseCost):
+    """
+    Jax-based Sum of Squared Error cost function.
+    """
+
+    def __init__(self, problem: BaseProblem):
+        super().__init__(problem)
+        self.model = self.problem.model
+        self.n_data = self.problem.n_data
+        if isinstance(self.model.solver, IDAKLUSolver):
+            self.model.jaxify_solver(t_eval=self.problem.domain_data)
+
+    def __call__(
+        self,
+        inputs: Inputs,
+        calculate_grad: bool = False,
+        apply_transform: bool = False,
+    ) -> Union[np.array, tuple[float, np.ndarray]]:
+        """
+        Computes the cost function for the given predictions.
+
+        Parameters
+        ----------
+        y : dict
+            The dictionary of predictions with keys designating the signals for fitting.
+        dy : np.ndarray, optional
+            The corresponding gradient with respect to the parameters for each signal.
+        calculate_grad : bool, optional
+            A bool condition designating whether to calculate the gradient.
+
+        Returns
+        -------
+        float
+            The Sum of Squared Error.
+        """
+        inputs = self.parameters.verify(inputs)
+        if calculate_grad != self.model.calculate_sensitivities:
+            self._update_solver_sensitivities(calculate_grad)
+
+        if calculate_grad:
+            y, dy = jax.value_and_grad(self.evaluate)(inputs)
+            return y, np.asarray(
+                list(dy.values())
+            )  # Convert grad to numpy for optimisers
+        else:
+            return np.asarray(self.evaluate(inputs))
+
+    def _update_solver_sensitivities(self, calculate_grad: bool) -> None:
+        """
+        Updates the solver's sensitivity calculation based on the gradient requirement.
+
+        Args:
+            calculate_grad (bool): Whether gradient calculation is required.
+        """
+
+        self.model.jaxify_solver(
+            t_eval=self.problem.domain_data, calculate_sensitivities=calculate_grad
+        )
+
+    @staticmethod
+    def check_sigma0(sigma0):
+        if not isinstance(sigma0, (int, float)) or sigma0 <= 0:
+            raise ValueError("sigma0 must be a positive number")
+        return float(sigma0)
+
+    def observed_fisher(self, inputs: Inputs):
+        """
+        Compute the observed fisher information matrix (FIM)
+        for the given inputs. This is done with the gradient
+        as the Hessian is not available.
+        """
+        _, grad = self.__call__(inputs, calculate_grad=True)
+        return jnp.square(grad) / self.n_data
+
+
+class JaxSumSquaredError(BaseJaxCost):
+    """
+    Jax-based Sum of Squared Error cost function.
+    """
+
+    def __init__(self, problem: BaseProblem):
+        super().__init__(problem)
+
+    def evaluate(self, inputs):
+        # Calculate residuals and error
+        y = self.problem.evaluate(inputs)
+        r = jnp.asarray([y[s] - self._target[s] for s in self.signal])
+        return jnp.sum(r**2)
+
+
+class JaxLogNormalLikelihood(BaseJaxCost, BaseLikelihood):
+    """
+    A Log-Normal Likelihood function. This function represents the
+    underlining observed data sampled from a Log-Normal distribution.
+
+    Parameters
+    -----------
+    problem: BaseProblem
+        The problem to fit of type `pybop.BaseProblem`
+    sigma0: float, optional
+        The variance in the measured data
+    """
+
+    def __init__(self, problem: BaseProblem, sigma0=0.02):
+        super().__init__(problem)
+        self.sigma = self.check_sigma0(sigma0)
+        self.sigma2 = jnp.square(self.sigma)
+        self._offset = 0.5 * self.n_data * jnp.log(2 * jnp.pi)
+        self._target_as_array = jnp.asarray([self._target[s] for s in self.signal])
+        self._log_target_sum = jnp.sum(jnp.log(self._target_as_array))
+        self._precompute()
+
+    def _precompute(self):
+        self._constant_term = (
+            -self._offset - self.n_data * jnp.log(self.sigma) - self._log_target_sum
+        )
+
+    def evaluate(self, inputs):
+        """
+        Evaluates the log-normal likelihood.
+        """
+        y = self.problem.evaluate(inputs)
+        e = jnp.asarray([jnp.log(y[s]) - jnp.log(self._target[s]) for s in self.signal])
+        likelihood = self._constant_term - jnp.sum(jnp.square(e)) / (2 * self.sigma2)
+        return likelihood
+
+
+class JaxGaussianLogLikelihoodKnownSigma(BaseJaxCost, BaseLikelihood):
+    """
+    A Jax implementation of the Gaussian Likelihood function.
+    This function represents the underlining observed data sampled
+    from a Gaussian distribution with known noise, `sigma0`.
+
+    Parameters
+    -----------
+    problem: BaseProblem
+        The problem to fit of type `pybop.BaseProblem`
+    sigma0: float, optional
+        The variance in the measured data
+    """
+
+    def __init__(self, problem: BaseProblem, sigma0=0.02):
+        super().__init__(problem)
+        self.sigma = self.check_sigma0(sigma0)
+        self.sigma2 = jnp.square(self.sigma)
+        self._offset = -0.5 * self.n_data * jnp.log(2 * jnp.pi * self.sigma2)
+        self._multip = -1 / (2.0 * self.sigma2)
+
+    def evaluate(self, inputs):
+        """
+        Evaluates the log-normal likelihood.
+        """
+        y = self.problem.evaluate(inputs)
+        e = jnp.asarray([y[s] - self._target[s] for s in self.signal])
+        likelihood = jnp.sum(self._offset + self._multip * jnp.sum(jnp.square(e)))
+        return likelihood