Adds support for IDAKLU w/ output variables

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#450](https://github.com/pybop-team/PyBOP/pull/450) - Adds support for IDAKLU with output variables, and corresponding examples, tests.
 - [#435](https://github.com/pybop-team/PyBOP/pull/435) - Adds SLF001 linting for private members.
 - [#418](https://github.com/pybop-team/PyBOP/issues/418) - Wraps the `get_parameter_info` method from PyBaMM to get a dictionary of parameter names and types.
 - [#413](https://github.com/pybop-team/PyBOP/pull/413) - Adds `DesignCost` functionality to `WeightedCost` class with additional tests.

examples/notebooks/solver_selection.ipynb

@@ -0,0 +1,245 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "expmkveO04pw"
+   },
+   "source": [
+    "## Investigate Different PyBaMM Solvers\n",
+    "\n",
+    "In this notebook, we discuss the process of changing PyBaMM solvers and the corresponding performance trade-offs with each. For further reading on different solvers, see the PyBaMM solver documentation:\n",
+    "\n",
+    "[[1]: PyBaMM Solvers](https://docs.pybamm.org/en/stable/source/api/solvers/index.html#)\n",
+    "\n",
+    "### Setting up the Environment\n",
+    "\n",
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP from its development branch and upgrade some dependencies:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "X87NUGPW04py",
+    "outputId": "0d785b07-7cff-4aeb-e60a-4ff5a669afbf"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Note: you may need to restart the kernel to use updated packages.\n",
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%pip install --upgrade pip ipywidgets -q\n",
+    "%pip install pybop -q\n",
+    "\n",
+    "import time\n",
+    "\n",
+    "import numpy as np\n",
+    "import pybamm\n",
+    "\n",
+    "import pybop"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "5XU-dMtU04p2"
+   },
+   "source": [
+    "### Setting up the model, and problem\n",
+    ".\n",
+    "We start by constructing a pybop model, and a synthetic dataset needed for the pybop problem we will be using for the solver benchmarking "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Model\n",
+    "parameter_set = pybop.ParameterSet.pybamm(\"Chen2020\")\n",
+    "model = pybop.lithium_ion.SPM(parameter_set=parameter_set)\n",
+    "\n",
+    "# Synthetic data\n",
+    "t_eval = np.arange(0, 900, 2)\n",
+    "values = model.predict(t_eval=t_eval)\n",
+    "\n",
+    "# Dataset\n",
+    "dataset = pybop.Dataset(\n",
+    "    {\n",
+    "        \"Time [s]\": t_eval,\n",
+    "        \"Current function [A]\": values[\"Current [A]\"].data,\n",
+    "        \"Voltage [V]\": values[\"Voltage [V]\"].data,\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "# Parameters\n",
+    "parameters = pybop.Parameters(\n",
+    "    pybop.Parameter(\n",
+    "        \"Negative electrode active material volume fraction\",\n",
+    "        prior=pybop.Gaussian(0.6, 0.02),\n",
+    "        bounds=[0.5, 0.8],\n",
+    "    ),\n",
+    "    pybop.Parameter(\n",
+    "        \"Positive electrode active material volume fraction\",\n",
+    "        prior=pybop.Gaussian(0.48, 0.02),\n",
+    "        bounds=[0.4, 0.7],\n",
+    "    ),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "n4OHa-aF04qA"
+   },
+   "source": [
+    "### Defining the solvers for benchmarking\n",
+    "\n",
+    "Now that we have set up the majority of the pybop objects, we construct the solvers we want to benchmark on the given model, and applied current."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "solvers = [\n",
+    "    pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6),\n",
+    "    pybamm.CasadiSolver(atol=1e-6, rtol=1e-6, mode=\"safe\"),\n",
+    "    pybamm.CasadiSolver(atol=1e-6, rtol=1e-6, mode=\"fast\"),\n",
+    "    pybamm.CasadiSolver(atol=1e-6, rtol=1e-6, mode=\"fast with events\"),\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, we construct a range of inputs for the parameters defined above, and select the number of instances in that range to benchmark on. For more statistically repeatable results, increase the variable `n` below."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n = 50  # Number of solves\n",
+    "inputs = list(zip(np.linspace(0.45, 0.6, n), np.linspace(0.45, 0.6, n)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, let's benchmark the solvers without sensitivities. This provides a reference for the non-gradient based pybop optimisers and samplers. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Time Evaluate IDA KLU solver: 0.521\n",
+      "Time Evaluate CasADi solver with 'safe' mode: 0.464\n",
+      "Time Evaluate CasADi solver with 'fast' mode: 0.458\n",
+      "Time Evaluate CasADi solver with 'fast with events' mode: 0.451\n"
+     ]
+    }
+   ],
+   "source": [
+    "for solver in solvers:\n",
+    "    model.solver = solver\n",
+    "    problem = pybop.FittingProblem(model, parameters, dataset)\n",
+    "\n",
+    "    start_time = time.time()\n",
+    "    for input_values in inputs:\n",
+    "        problem.evaluate(inputs=input_values)\n",
+    "    print(f\"Time Evaluate {solver.name}: {time.time() - start_time:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Excellent, given the above results, we know which solver we should select for optimisation on your machine, i.e. the one with the smallest time. \n",
+    "\n",
+    "Next, let's repeat the same toy problem, but for the gradient-based cost evaluation,"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Time EvaluateS1 IDA KLU solver: 1.478\n",
+      "Time EvaluateS1 CasADi solver with 'safe' mode: 2.836\n",
+      "Time EvaluateS1 CasADi solver with 'fast' mode: 2.699\n",
+      "Time EvaluateS1 CasADi solver with 'fast with events' mode: 2.714\n"
+     ]
+    }
+   ],
+   "source": [
+    "for solver in solvers:\n",
+    "    model.solver = solver\n",
+    "    problem = pybop.FittingProblem(model, parameters, dataset)\n",
+    "\n",
+    "    start_time = time.time()\n",
+    "    for input_values in inputs:\n",
+    "        problem.evaluateS1(inputs=input_values)\n",
+    "    print(f\"Time EvaluateS1 {solver.name}: {time.time() - start_time:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we have the relevant information for the gradient-based optimisers. Likewise to the above results, we should select the solver with the smallest time."
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

examples/scripts/selecting_a_solver.py

@@ -0,0 +1,60 @@
+import time
+
+import numpy as np
+import pybamm
+
+import pybop
+
+# Parameter set and model definition
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+model = pybop.lithium_ion.SPM(parameter_set=parameter_set)
+
+solvers = [
+    pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6),
+    pybamm.CasadiSolver(mode="safe", atol=1e-6, rtol=1e-6),
+    pybamm.CasadiSolver(mode="fast with events", atol=1e-6, rtol=1e-6),
+]
+
+# 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
+experiment = pybop.Experiment([("Discharge at 0.5C for 10 minutes (3 second period)")])
+values = model.predict(
+    initial_state={"Initial open-circuit voltage [V]": 4.2}, 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,
+    }
+)
+
+# Create the list of input dicts
+n = 150  # Number of solves
+inputs = list(zip(np.linspace(0.45, 0.6, n), np.linspace(0.45, 0.6, n)))
+
+# Iterate over the solvers and print benchmarks
+for solver in solvers:
+    model.solver = solver
+    problem = pybop.FittingProblem(model, parameters, dataset)
+
+    start_time = time.time()
+    for input_values in inputs:
+        problem.evaluate(inputs=input_values)
+    print(f"Time Evaluate {solver.name}: {time.time() - start_time:.3f}")
+
+    start_time = time.time()
+    for input_values in inputs:
+        problem.evaluateS1(inputs=input_values)
+    print(f"Time EvaluateS1 {solver.name}: {time.time() - start_time:.3f}")

pybop/models/base_model.py

@@ -803,3 +803,7 @@ def spatial_methods(self):
     @property
     def solver(self):
         return self._solver
+
+    @solver.setter
+    def solver(self, solver):
+        self._solver = solver.copy() if solver is not None else None

pybop/problems/base_problem.py

@@ -1,5 +1,8 @@
 from typing import Optional
 
+import numpy as np
+from pybamm import IDAKLUSolver
+
 from pybop import BaseModel, Dataset, Parameter, Parameters
 from pybop.parameters.parameter import Inputs
 
@@ -68,6 +71,21 @@ def __init__(
         self._time_data = None
         self._target = None
         self.verbose = False
+        self.failure_output = np.asarray([np.inf])
+
+        # If model.solver is IDAKLU, set output vars for improved performance
+        self.output_vars = tuple(self.signal + self.additional_variables)
+        if self._model is not None and isinstance(self._model.solver, IDAKLUSolver):
+            self._solver_copy = self._model.solver.copy()
+            self._model.solver = IDAKLUSolver(
+                atol=self._solver_copy.atol,
+                rtol=self._solver_copy.rtol,
+                root_method=self._solver_copy.root_method,
+                root_tol=self._solver_copy.root_tol,
+                extrap_tol=self._solver_copy.extrap_tol,
+                options=self._solver_copy._options,  # noqa: SLF001
+                output_variables=self.output_vars,
+            )
 
     @property
     def n_parameters(self):

pybop/problems/fitting_problem.py

@@ -107,7 +107,7 @@ def evaluate(
         except Exception as e:
             if self.verbose:
                 print(f"Simulation error: {e}")
-            return {signal: [np.inf] for signal in self.signal}
+            return {signal: self.failure_output for signal in self.signal}
 
         return {
             signal: sol[signal].data
@@ -137,7 +137,9 @@ def evaluateS1(self, inputs: Inputs):
             )
         except Exception as e:
             print(f"Error: {e}")
-            return {signal: [np.inf] for signal in self.signal}, [np.inf]
+            return {
+                signal: self.failure_output for signal in self.signal
+            }, self.failure_output
 
         y = {signal: sol[signal].data for signal in self.signal}