{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cycle ageing example\n",
    "\n",
    "This notebook demonstrates how to use the `CycleAgeing` objective function to fit the\n",
    "SEI solvent diffusivity in the \"solvent-diffusion limited\" SEI model using synthetic cycling data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ionworkspipeline as iwp\n",
    "from ionworkspipeline.data_fits.models.metrics import First, Last\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pybamm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin by generating some synthetic cycling data. We choose the SPMe model and the \"solvent-diffusion limited\" SEI model and simulate CCCV charge/discharge cycling with RPTs every 10 cycles. For the data, we record the cycle number, the loss of lithium inventory, and the C/5 discharge capacity. We choose the parameter values so that we get a large loss of lithium inventory even after only 50 cycles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model and parameters\n",
    "model = pybamm.lithium_ion.SPMe({\"SEI\": \"solvent-diffusion limited\"})\n",
    "parameter_values = pybamm.ParameterValues(\"Chen2020\")\n",
    "parameter_values.update({\"SEI solvent diffusivity [m2.s-1]\": 5e-19})\n",
    "\n",
    "# Define RPT and CCCV cycles\n",
    "rpt_cycle = (\n",
    "    \"Discharge at C/5 until 2.5 V (60s period)\",\n",
    "    \"Rest for 1 hour\",\n",
    "    \"Charge at C/5 until 4.2 V\",\n",
    "    \"Hold at 4.2 V until 10 mA\",\n",
    "    \"Rest for 1 hour\",\n",
    ")\n",
    "cccv_cycle = (\n",
    "    \"Discharge at 1C until 2.5 V\",\n",
    "    \"Rest for 1 hour\",\n",
    "    \"Charge at C/2 until 4.2 V\",\n",
    "    \"Hold at 4.2 V until 10 mA\",\n",
    "    \"Rest for 1 hour\",\n",
    ")\n",
    "# One block = 10 CCCV cycles + 1 RPT\n",
    "one_block = [cccv_cycle] * 10 + [rpt_cycle]\n",
    "# Simulate initial RPT, then 10 CCCV cycles + 1 RPT repeated 5 times\n",
    "cycles_list = [rpt_cycle] + one_block * 5\n",
    "experiment = pybamm.Experiment(cycles_list)\n",
    "sim = pybamm.Simulation(model, experiment=experiment, parameter_values=parameter_values)\n",
    "sol = sim.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can extract the synthetic data from the simulation and plot the RPTs to observe the effect of the degradation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Indices of cycles that are RPTs (0, 11, 21, ..., as RPT cycles are every 11 cycles)\n",
    "rpt_idxs = np.arange(0, len(cycles_list), 11)\n",
    "\n",
    "# Extract the data and plot the RPTs\n",
    "data_rows = []\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].set_xlabel(\"Time [h]\")\n",
    "ax[0].set_ylabel(\"Voltage [V]\")\n",
    "ax[1].set_xlabel(\"Cycle number\")\n",
    "ax[1].set_ylabel(\"LLI [%]\")\n",
    "colors = plt.cm.viridis(np.linspace(0, 1, len(rpt_idxs)))\n",
    "for idx, color in zip(rpt_idxs, colors, strict=False):\n",
    "    # Extract the cycle solution and the C/5 discharge step\n",
    "    cycle_sol = sol.cycles[idx]\n",
    "    step_sol = cycle_sol.steps[0]\n",
    "\n",
    "    # Store LLI\n",
    "    LLI = cycle_sol[\"LLI [%]\"].data[0]\n",
    "    # Calculate C/5 discharge capacity\n",
    "    q = step_sol[\"Discharge capacity [A.h]\"].data\n",
    "    q_c5 = q[-1] - q[0]\n",
    "    data_rows.append(\n",
    "        {\n",
    "            \"Cycle number\": idx,\n",
    "            \"LLI [%]\": LLI,\n",
    "            \"C/5 capacity [A.h]\": q_c5,\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # Plot C/5 discharge from RPT and LLI\n",
    "    time = step_sol[\"Time [h]\"].data\n",
    "    time -= time[0]\n",
    "    voltage = step_sol[\"Voltage [V]\"].data\n",
    "    ax[0].plot(time, voltage, color=color, label=f\"Cycle {idx}\")\n",
    "    ax[1].scatter(idx, LLI, color=color)\n",
    "\n",
    "ax[0].legend()\n",
    "\n",
    "# Convert to DataFrame\n",
    "data = pd.DataFrame(data_rows)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we create the objective function. We pass the data and the options to the `CycleAgeing` objective function. The options include the model, experiment, the objective variables (i.e. the variables we want to fit to), and the metrics that define how to extract values from the simulation.\n",
    "\n",
    "The `metrics` option accepts a dictionary mapping variable names to `Metric` objects. We can combine metrics to create new metrics, which enables calculating more complex quantities that aren't directly in the model, such as the C/5 discharge capacity or the pulse resistance. Here we use `First(\"LLI [%]\").by_cycle()` to extract the first value of the loss of lithium inventory at each cycle and use the `Last(\"Discharge capacity [A.h]\").by_cycle(step=0) - First(\"Discharge capacity [A.h]\").by_cycle(step=0)` to extract the C/5 discharge capacity at each cycle. The `.by_cycle()` method evaluates the metric for each cycle in the simulation. We pass the `step=0` argument to the `by_cycle()` method to extract the values from the first step of the cycle.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = iwp.objectives.CycleAgeing(\n",
    "    data,\n",
    "    options={\n",
    "        \"model\": model,\n",
    "        \"experiment\": experiment,\n",
    "        \"objective variables\": [\"LLI [%]\", \"C/5 capacity [A.h]\"],\n",
    "        \"metrics\": {\n",
    "            \"LLI [%]\": First(\"LLI [%]\").by_cycle(),\n",
    "            \"C/5 capacity [A.h]\": Last(\"Discharge capacity [A.h]\").by_cycle(step=0)\n",
    "            - First(\"Discharge capacity [A.h]\").by_cycle(step=0),\n",
    "        },\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then create the data fit. We pass the objective function and the parameters to the `DataFit` class. The parameters include the initial values and bounds for the parameters we want to fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"SEI solvent diffusivity [m2.s-1]\": iwp.Parameter(\n",
    "        \"SEI solvent diffusivity [m2.s-1]\",\n",
    "        bounds=(1e-20, 1e-18),\n",
    "    )\n",
    "}\n",
    "optimizer = iwp.optimizers.Pints(method=\"PSO\", log_to_screen=True)\n",
    "data_fit = iwp.DataFit(objective, parameters=parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we run the data fit. We pass the known parameters (i.e. those we are not fitting) to the data fit. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "known_parameters = {k: v for k, v in parameter_values.items() if k not in parameters}\n",
    "data_fit.run(known_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we plot the fit results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = data_fit.plot_fit_results()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fit.plot_trace()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}