{
 "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": [
      "/Users/engs2510/Documents/Git/PyBOP/.nox/notebooks-overwrite/bin/python3: No module named pip\r\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n",
      "/Users/engs2510/Documents/Git/PyBOP/.nox/notebooks-overwrite/bin/python3: No module named pip\r\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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\n",
    "\n",
    "pybop.PlotlyManager().pio.renderers.default = \"notebook_connected\""
   ]
  },
  {
   "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.303\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time Evaluate CasADi solver with 'safe' mode: 1.076\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time Evaluate CasADi solver with 'fast' mode: 1.003\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time Evaluate CasADi solver with 'fast with events' mode: 1.008\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.671\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time EvaluateS1 CasADi solver with 'safe' mode: 4.479\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time EvaluateS1 CasADi solver with 'fast' mode: 3.440\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time EvaluateS1 CasADi solver with 'fast with events' mode: 3.287\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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}