{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting OCP hysteresis\n",
    "\n",
    "In this example, we show how to fit the lithiation and delithiation OCPs at the same time in a way that is consistent with the MSMR model. In particular, we want the total capacity to be the same during lithiation and delithiation. To ensure this, we fit a single capacity for each species, but allow the reaction parameters $X_j$, $U0_j$ and $\\omega_j$ to be different for lithiation and delithiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ionworkspipeline as iwp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.signal import savgol_filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use data from half-cells constructed from an LGM50 cell, as reported in https://github.com/paramm-team/pybamm-param."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {}\n",
    "Q_max_data = 0\n",
    "fig, ax = plt.subplots(1, 3, figsize=(10, 5))\n",
    "for direction in [\"lithiation\", \"delithiation\"]:\n",
    "    # Smooth\n",
    "    raw_ocp_data = pd.read_csv(f\"anode_ocp_{direction}.csv\", comment=\"#\")\n",
    "    q_smoothed = savgol_filter(\n",
    "        raw_ocp_data[\"Capacity [A.h]\"].values, window_length=11, polyorder=1\n",
    "    )\n",
    "    U_smoothed = savgol_filter(\n",
    "        raw_ocp_data[\"Voltage [V]\"].values, window_length=11, polyorder=1\n",
    "    )\n",
    "\n",
    "    # Store in a dataframe\n",
    "    ocp_data = pd.DataFrame({\"Capacity [A.h]\": q_smoothed, \"Voltage [V]\": U_smoothed})\n",
    "    # Preprocess\n",
    "    ocp_data = iwp.data_fits.preprocess.sort_capacity_and_ocp(ocp_data)\n",
    "    ocp_data = iwp.data_fits.preprocess.remove_ocp_extremes(ocp_data)\n",
    "    q = ocp_data[\"Capacity [A.h]\"].values\n",
    "    U = ocp_data[\"Voltage [V]\"].values\n",
    "    data[direction] = ocp_data\n",
    "    # Maximum capacity, will be used later as an initial guess for the electrode capacity\n",
    "    Q_max_data = max(Q_max_data, ocp_data[\"Capacity [A.h]\"].max())\n",
    "\n",
    "    # Plot\n",
    "    dUdQ = -np.gradient(U, q)\n",
    "    dQdU = -np.gradient(q, U)\n",
    "    ax[0].plot(q, U, \"-x\", label=direction)\n",
    "    ax[0].set_xlabel(\"Capacity [A.h]\")\n",
    "    ax[0].set_ylabel(\"Voltage [V]\")\n",
    "    ax[1].plot(q, dUdQ, \"-x\")\n",
    "    ax[1].set_xlabel(\"Capacity [A.h]\")\n",
    "    ax[1].set_ylabel(\"dU/dQ [V/A.h]\")\n",
    "    ax[1].set_ylim(0, 0.5)\n",
    "    ax[2].plot(dQdU, U, \"-x\")\n",
    "    ax[2].set_xlabel(\"dQ/dU [A.h/V]\")\n",
    "    ax[2].set_ylabel(\"Voltage [V]\")\n",
    "    ax[2].set_xlim(0, 200)\n",
    "ax[0].legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we set up the `Parameter` objects for the fit. For the reaction parameters, we use the values from [Verbrugge et al 2017 J. Electrochem. Soc. 164 E3243](https://iopscience.iop.org/article/10.1149/2.0341708jes) as an initial guess. We can access the values from the `Material` object, which is part of our library of parameter values.\n",
    "\n",
    "We set up a bounds function that returns the bounds for the reaction parameters. Here we use the same bounds for lithiation and delithiation, and for each reaction index $j$. We could construct a more complex bounds function that allows the bounds to be different for lithiation and delithiation, and/or different for each reaction index $j$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial guess of Graphite from Verbrugge et al 2017 J. Electrochem. Soc. 164 E3243\n",
    "msmr_param_init = iwp.Material(\"GRAPHITE_VERBRUGGE2017\").parameter_values\n",
    "\n",
    "\n",
    "# Define bounds for the MSMR parameters\n",
    "def bounds_function(var, initial_value, Q_total):\n",
    "    if var.startswith(\"X\"):\n",
    "        return 0, 1\n",
    "    elif var.startswith(\"U\"):\n",
    "        return max(0, initial_value - 0.2), initial_value + 0.2\n",
    "    elif var.startswith(\"w\"):\n",
    "        return 1e-2, 100\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown parameter {var}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create a dictionary of parameters to fit. We fit a single capacity for each species, but allow the lower excess capacity to be different for lithiation and delithiation, since the data in both directions may not span the same range of capacities. For the reaction parameters, we use the the helper function `get_msmr_params_for_fit` to create the `Parameter` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_to_fit = {\n",
    "    \"Negative electrode capacity [A.h]\": iwp.Parameter(\n",
    "        \"Q_n\", initial_value=Q_max_data, bounds=(Q_max_data * 0.8, Q_max_data * 1.2)\n",
    "    ),\n",
    "    \"Negative electrode lower excess capacity (delithiation) [A.h]\": iwp.Parameter(\n",
    "        \"Q_n_d_lowex\", initial_value=0.01 * Q_max_data, bounds=(0, 0.1 * Q_max_data)\n",
    "    ),\n",
    "    \"Negative electrode lower excess capacity (lithiation) [A.h]\": iwp.Parameter(\n",
    "        \"Q_n_l_lowex\", initial_value=0.01 * Q_max_data, bounds=(0, 0.1 * Q_max_data)\n",
    "    ),\n",
    "}\n",
    "params_to_fit.update(\n",
    "    {\n",
    "        **iwp.objectives.get_msmr_params_for_fit(\n",
    "            msmr_param_init,\n",
    "            \"negative\",\n",
    "            Q_max_data,\n",
    "            bounds_function=bounds_function,\n",
    "            species_format=\"Xj\",\n",
    "            hysteresis=True,\n",
    "        )\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we create the objectives. For this fit, we use the `OpenCircuitHalfCellHysteresisCallback` that shows both the fit results on a single plot. We create a single instance of the callback and pass it to each objective."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objectives = {}\n",
    "callback = iwp.objectives.OpenCircuitHalfCellHysteresisCallback()\n",
    "for direction in [\"lithiation\", \"delithiation\"]:\n",
    "    objectives[direction] = iwp.objectives.MSMRHalfCell(\n",
    "        data[direction],\n",
    "        options={\n",
    "            \"model\": iwp.data_fits.models.MSMRHalfCellModel(\n",
    "                \"negative\", options={\"species format\": \"Xj\", \"direction\": direction}\n",
    "            ),\n",
    "            \"dUdQ cutoff\": 0.2,\n",
    "            \"voltage limits\": (0, 1),\n",
    "        },\n",
    "        callbacks=callback,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we create the `DataFit` object, run the fit,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fit = iwp.DataFit(objectives, parameters=params_to_fit)\n",
    "\n",
    "results = data_fit.run({\"Ambient temperature [K]\": 298.15})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = callback.plot_fit_results()\n",
    "_ = axes[2].set_xlim(0, 200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}