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.
 - [#364](https://github.com/pybop-team/PyBOP/pull/364) - Adds the MultiFittingProblem class and the multi_fitting example script.
 - [#444](https://github.com/pybop-team/PyBOP/issues/444) - Merge `BaseModel` `build()` and `rebuild()` functionality.
 - [#435](https://github.com/pybop-team/PyBOP/pull/435) - Adds SLF001 linting for private members.

examples/notebooks/equivalent_circuit_identification.ipynb

@@ -11,7 +11,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

examples/notebooks/multi_model_identification.ipynb

@@ -12,7 +12,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

examples/notebooks/multi_optimiser_identification.ipynb

@@ -12,7 +12,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

examples/notebooks/optimiser_calibration.ipynb

@@ -12,7 +12,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

examples/notebooks/pouch_cell_identification.ipynb

@@ -12,7 +12,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

examples/notebooks/solver_selection.ipynb

@@ -0,0 +1,261 @@
+{
+ "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 and upgrade 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": {},
+   "source": [
+    "Let's fix the random seed in order to generate consistent output during development, although this does not need to be done in practice."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.random.seed(8)"
+   ]
+  },
+  {
+   "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.477\n",
+      "Time Evaluate CasADi solver with 'safe' mode: 0.464\n",
+      "Time Evaluate CasADi solver with 'fast' mode: 0.446\n",
+      "Time Evaluate CasADi solver with 'fast with events' mode: 0.443\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: 0.997\n",
+      "Time EvaluateS1 CasADi solver with 'safe' mode: 2.743\n",
+      "Time EvaluateS1 CasADi solver with 'fast' mode: 2.689\n",
+      "Time EvaluateS1 CasADi solver with 'fast with events' mode: 2.688\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/notebooks/spm_AdamW.ipynb

@@ -16,7 +16,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

examples/notebooks/spm_electrode_design.ipynb

@@ -18,7 +18,7 @@
     "\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:"
+    "Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:"
    ]
   },
   {

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}")