{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Particle Swarm Optimizer\n",
    "\n",
    "In this example, we will use the Particle Swarm Optimizer (PSO) to fit a battery model to some synthetic data.\n",
    "\n",
    "We begin by generating some synthetic data from a battery model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pybamm\n",
    "import ionworkspipeline as iwp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# Load the model and parameter values\n",
    "model = pybamm.lithium_ion.SPM()\n",
    "parameter_values = iwp.ParameterValues(\"Chen2020\")\n",
    "parameter_values.update(\n",
    "    {\n",
    "        \"Negative electrode active material volume fraction\": 0.75,\n",
    "        \"Positive electrode active material volume fraction\": 0.665,\n",
    "    }\n",
    ")\n",
    "\n",
    "# Generate the data\n",
    "t = np.arange(0, 900, 3)\n",
    "sim = iwp.Simulation(\n",
    "    model, parameter_values=parameter_values, t_eval=[0, 900], t_interp=t\n",
    ")\n",
    "sol = sim.solve()\n",
    "data = pd.DataFrame(\n",
    "    {x: sim.solution[x](t) for x in [\"Time [s]\", \"Current [A]\", \"Voltage [V]\"]}\n",
    ")\n",
    "# add noise to the data\n",
    "sigma = 0.001\n",
    "data[\"Voltage [V]\"] += np.random.normal(0, sigma, len(data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define the parameters to fit, along with initial guesses and bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"Negative electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Negative electrode active material volume fraction\",\n",
    "        initial_value=0.5,\n",
    "        bounds=(0, 1),\n",
    "    ),\n",
    "    \"Positive electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Positive electrode active material volume fraction\",\n",
    "        initial_value=0.5,\n",
    "        bounds=(0, 1),\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set up the objective function, which in this case is the current driven objective. This takes the time vs current data and runs the model, comparing the model's predictions for the voltage with the experimental data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = iwp.objectives.CurrentDriven(data, options={\"model\": model})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we set up the `DataFit` object, which takes the objective function, the parameters to fit, and the optimizer as inputs. Here we use the Particle Swarm Optimizer (PSO) from the `pints` library and set the maximum number of iterations to 100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = iwp.optimizers.Pints(max_iterations=100, method=\"PSO\")\n",
    "\n",
    "data_fit = iwp.DataFit(objective, parameters=parameters, optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we run the data fit, passing the parameters that are not being fit as a dictionary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_for_pipeline = {k: v for k, v in parameter_values.items() if k not in parameters}\n",
    "results = data_fit.run(params_for_pipeline)\n",
    "for k, v in results.items():\n",
    "    print(f\"{k}: {parameter_values[k]:.3e} (true) {v:.3e} (fit)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we plot the results and the trace of the cost function and parameter values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = data_fit.plot_fit_results()\n",
    "_ = data_fit.plot_trace()"
   ]
  }
 ],
 "metadata": {
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}