{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Half-cell OCP fitting\n",
    "\n",
    "In this example, we show how to fit the half-cell Open-Circuit Potential to pseudo-OCP data. Our recommended workflow for fitting OCPs is to first fit the MSMR model to the data and then create an interpolant to find a function $U(x)$, where $x$ is the stoichiometry. For comparison, we also show how to fit a function $U(x)$ directly to pseudo-OCP data."
   ]
  },
  {
   "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 pybamm\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.signal import savgol_filter"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-process and visualize the data\n",
    "First, we load and pre-process the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data\n",
    "raw_ocp_data = pd.read_csv(\"anode_ocp_delithiation.csv\", comment=\"#\")\n",
    "q = raw_ocp_data[\"Capacity [A.h]\"].values\n",
    "U = raw_ocp_data[\"Voltage [V]\"].values\n",
    "\n",
    "# Smooth\n",
    "q_smoothed = savgol_filter(q, window_length=11, polyorder=1)\n",
    "U_smoothed = savgol_filter(U, window_length=11, polyorder=1)\n",
    "dQdU_smoothed = -np.gradient(q_smoothed, U_smoothed)\n",
    "\n",
    "# Store in a dataframe\n",
    "ocp_data = pd.DataFrame({\"Capacity [A.h]\": q_smoothed, \"Voltage [V]\": U_smoothed})\n",
    "ocp_data = iwp.data_fits.preprocess.sort_capacity_and_ocp(ocp_data)\n",
    "# Maximum capacity, will be used later as an initial guess for the electrode capacity\n",
    "Q_max_data = ocp_data[\"Capacity [A.h]\"].max()\n",
    "\n",
    "# Plot\n",
    "_, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].plot(q_smoothed, U_smoothed, \"-x\")\n",
    "ax[0].set_xlabel(\"Capacity [A.h]\")\n",
    "ax[0].set_ylabel(\"Voltage [V]\")\n",
    "ax[1].plot(dQdU_smoothed, U_smoothed, \"-x\")\n",
    "ax[1].set_xlabel(\"dQ/dU [A.h/V]\")\n",
    "ax[1].set_ylabel(\"Voltage [V]\")\n",
    "ax[1].set_xlim(0, 200)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit the MSMR model to data\n",
    "\n",
    "We begin by setting up initial guesses for our MSMR parameters."
   ]
  },
  {
   "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 = {\n",
    "    \"U0_p_0\": 0.088,\n",
    "    \"X_p_0\": 0.43,\n",
    "    \"w_p_0\": 0.086,\n",
    "    \"U0_p_1\": 0.13,\n",
    "    \"X_p_1\": 0.24,\n",
    "    \"w_p_1\": 0.080,\n",
    "    \"U0_p_2\": 0.14,\n",
    "    \"X_p_2\": 0.15,\n",
    "    \"w_p_2\": 0.72,\n",
    "    \"U0_p_3\": 0.17,\n",
    "    \"X_p_3\": 0.055,\n",
    "    \"w_p_3\": 2.5,\n",
    "    \"U0_p_4\": 0.21,\n",
    "    \"X_p_4\": 0.067,\n",
    "    \"w_p_4\": 0.095,\n",
    "    \"U0_p_5\": 0.36,\n",
    "    \"X_p_5\": 0.055,\n",
    "    \"w_p_5\": 6.0,\n",
    "    \"Number of reactions in positive electrode\": 6,\n",
    "    \"Positive electrode lower excess capacity [A.h]\": 0.0,\n",
    "    \"Positive electrode capacity [A.h]\": Q_max_data,\n",
    "    \"Ambient temperature [K]\": 298.15,\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot these against the data and, if required, manually adjust them to get a better starting point. For this example, let's keep the Graphite parameters as our initial guess."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = iwp.objectives.plot_each_species_msmr(\n",
    "    ocp_data, msmr_param_init, \"positive\", voltage_limits=(0, 1.5)\n",
    ")\n",
    "ax[1].set_xlim(0, 50)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can set up the objective and parameters to fit, run the data fit and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit the electrode capacity and lower excess capacity by making them Parameter objects\n",
    "msmr_params = {\n",
    "    \"Positive electrode capacity [A.h]\": iwp.Parameter(\n",
    "        \"Q\", initial_value=Q_max_data, bounds=(0.9 * Q_max_data, 1.2 * Q_max_data)\n",
    "    ),\n",
    "    \"Positive electrode lower excess capacity [A.h]\": (\n",
    "        iwp.Parameter(\n",
    "            \"Q_lowex\", initial_value=0.01 * Q_max_data, bounds=(0, 0.01 * Q_max_data)\n",
    "        )\n",
    "    ),\n",
    "}\n",
    "# Add Parameter objects for the MSMR parameters X_j, U0_j, and w_j\n",
    "for k, v in msmr_param_init.items():\n",
    "    if \"X\" in k:\n",
    "        bounds = (0, 1)\n",
    "    elif \"U0_\" in k:\n",
    "        # allow a 500 mV variation either way\n",
    "        bounds = (max(0, v - 0.5), v + 0.5)\n",
    "    elif \"w_\" in k:\n",
    "        bounds = (1e-2, 100)\n",
    "    else:\n",
    "        continue\n",
    "    msmr_params[k] = iwp.Parameter(k, initial_value=v, bounds=bounds)\n",
    "\n",
    "# Calculate dUdQ cutoff to prevent overfitting to extreme dUdQ values\n",
    "dUdq_cutoff = iwp.data_fits.util.calculate_dUdQ_cutoff(ocp_data)\n",
    "\n",
    "# Set up model, objective and fit\n",
    "model = iwp.data_fits.models.MSMRHalfCellModel(\n",
    "    \"positive\", options={\"species format\": \"Xj\"}\n",
    ")\n",
    "objective = iwp.data_fits.objectives.MSMRHalfCell(\n",
    "    ocp_data, options={\"model\": model, \"dUdQ cutoff\": dUdq_cutoff}\n",
    ")\n",
    "ocp_msmr = iwp.DataFit(objective, parameters=msmr_params)\n",
    "\n",
    "# Run the fit\n",
    "results = ocp_msmr.run({\"Ambient temperature [K]\": 298.15})\n",
    "\n",
    "# Plot results\n",
    "fig_ax = ocp_msmr.plot_fit_results()\n",
    "axes = fig_ax[\"MSMRHalfCell\"][0][1]\n",
    "_ = axes[2].set_xlim(0, 100)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we can look at the fitted parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.callbacks[\"MSMRHalfCell\"].callbacks[0].fit_results_[\"outputs\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can access tabulated capacity vs OCP data from the MSMR callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = results.callbacks[\"MSMRHalfCell\"].callbacks[0].fit_results_[\"outputs\"]\n",
    "q = outputs[\"Capacity [A.h]\"]\n",
    "U = outputs[\"Voltage [V]\"]\n",
    "ocp_fit_data = pd.DataFrame({\"Capacity [A.h]\": q, \"Voltage [V]\": U})\n",
    "ocp_fit_data.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit OCP as a function of stoichiometry\n",
    "\n",
    "Now, we directly fit a function $U(x)$ to the pseudo-OCP data.\n",
    "\n",
    "First, we define the function and initial guesses for the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def U_half_cell(theta):\n",
    "    u00 = pybamm.Parameter(\"u00_pos\")\n",
    "    u10 = pybamm.Parameter(\"u10_pos\")\n",
    "    u11 = pybamm.Parameter(\"u11_pos\")\n",
    "    u20 = pybamm.Parameter(\"u20_pos\")\n",
    "    u21 = pybamm.Parameter(\"u21_pos\")\n",
    "    u22 = pybamm.Parameter(\"u22_pos\")\n",
    "    u30 = pybamm.Parameter(\"u30_pos\")\n",
    "    u31 = pybamm.Parameter(\"u31_pos\")\n",
    "    u32 = pybamm.Parameter(\"u32_pos\")\n",
    "    u40 = pybamm.Parameter(\"u40_pos\")\n",
    "    u41 = pybamm.Parameter(\"u41_pos\")\n",
    "    u42 = pybamm.Parameter(\"u42_pos\")\n",
    "\n",
    "    u_eq = (\n",
    "        u00\n",
    "        + u10 * pybamm.exp(-u11 * theta)\n",
    "        + u20 * pybamm.tanh(u21 * (theta - u22))\n",
    "        + u30 * pybamm.tanh(u31 * (theta - u32))\n",
    "        + u40 * pybamm.tanh(u41 * (theta - u42))\n",
    "    )\n",
    "\n",
    "    return u_eq\n",
    "\n",
    "\n",
    "ocp_param_init = {\n",
    "    \"u00_pos\": 0.2,\n",
    "    \"u10_pos\": 2,\n",
    "    \"u11_pos\": 30,\n",
    "    \"u20_pos\": -0.1,\n",
    "    \"u21_pos\": 30,\n",
    "    \"u22_pos\": 0.1,\n",
    "    \"u30_pos\": -0.1,\n",
    "    \"u31_pos\": 30,\n",
    "    \"u32_pos\": 0.3,\n",
    "    \"u40_pos\": -0.1,\n",
    "    \"u41_pos\": 30,\n",
    "    \"u42_pos\": 0.6,\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we set up our `DataFit` object and fit the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up Parameter objects to fit\n",
    "ocp_params = {k: iwp.Parameter(k, initial_value=v) for k, v in ocp_param_init.items()}\n",
    "# Add the functional form of the OCP\n",
    "ocp_params = {\"Positive electrode OCP [V]\": U_half_cell, **ocp_params}\n",
    "\n",
    "# Use the experimental capacity to map between capacity in the experiment and the\n",
    "# stoichiometry. In practice we would calculate the capacity from other means,\n",
    "# (e.g. electrode loading) or use the MSMR model to fit the electrode capacity.\n",
    "known_parameter_values = {\"Positive electrode capacity [A.h]\": Q_max_data}\n",
    "\n",
    "# Set up objective and fit\n",
    "objective = iwp.objectives.OCPHalfCell(\n",
    "    \"positive\",\n",
    "    ocp_data,\n",
    ")\n",
    "ocp = iwp.DataFit(objective, parameters=ocp_params)\n",
    "\n",
    "# Fit\n",
    "params_fit_ocp = ocp.run(known_parameter_values)\n",
    "\n",
    "# Plot\n",
    "_ = ocp.plot_fit_results()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we look at the fitted parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_fit_ocp"
   ]
  },
  {
   "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
}