{
 "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 getting literature values for the MSMR parameters from the ionworkspipeline library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lib = iwp.Library()\n",
    "gr = lib[\"Graphite - Verbrugge 2017\"]\n",
    "gr_params = gr.parameter_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This gives us initial guesses for the MSMR species parameters ($X_j$, $U^0_j$, and $w_j$). We also need estimates for the total electrode capacity and lower excess capacity. We can use the helper function `get_initial_capacity_and_lower_excess_capacity` to estimate these from the species parameters and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gr_params.update(\n",
    "    iwp.data_fits.objectives.get_initial_capacity_and_lower_excess_capacity(\n",
    "        gr_params, \"negative\", ocp_data\n",
    "    )\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, gr_params, \"negative\", 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": [
    "# Loop over the initial guesses and set up the parameters to fit with appropriate bounds\n",
    "paramaters_to_fit = {}\n",
    "for k, v in gr_params.items():\n",
    "    if \"X\" in k:\n",
    "        bounds = (0, 1)\n",
    "    elif \"U0_\" in k:\n",
    "        bounds = (max(0, v - 0.5), v + 0.5)\n",
    "    elif \"w_\" in k:\n",
    "        bounds = (1e-2, 100)\n",
    "    elif \"lower excess\" in k:\n",
    "        bounds = (0, 0.1 * Q_max_data)\n",
    "    elif \"capacity\" in k:\n",
    "        bounds = (Q_max_data, 1.5 * Q_max_data)\n",
    "    else:\n",
    "        continue\n",
    "    paramaters_to_fit[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",
    "    \"negative\", 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=paramaters_to_fit)\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": "env",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}