{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting diffusivity \n",
    "\n",
    "In this example, we showcase some common workflows for fitting diffusion coefficients using pulse data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pybamm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import ionworkspipeline as iwp\n",
    "from scipy.integrate import cumulative_trapezoid\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Generate synthetic data\n",
    "\n",
    "For this example, we set up a Single Particle Model with a stoichiometry dependent diffusivity. We first set up the model and parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "options = {\"working electrode\": \"positive\"}\n",
    "model = pybamm.lithium_ion.SPM(options)\n",
    "parameter_values = iwp.ParameterValues(\"Xu2019\")\n",
    "\n",
    "\n",
    "def D_p(sto, T):\n",
    "    \"NMC diffusivity from O'Regan et al. 2021\"\n",
    "    a1 = -0.9231\n",
    "    a2 = -0.4066\n",
    "    a3 = -0.993\n",
    "    b1 = 0.3216\n",
    "    b2 = 0.4532\n",
    "    b3 = 0.8098\n",
    "    c0 = -13.96\n",
    "    c1 = 0.002534\n",
    "    c2 = 0.003926\n",
    "    c3 = 0.09924\n",
    "\n",
    "    D = (\n",
    "        10\n",
    "        ** (\n",
    "            c0\n",
    "            + a1 * np.exp(-((sto - b1) ** 2) / c1)\n",
    "            + a2 * np.exp(-((sto - b2) ** 2) / c2)\n",
    "            + a3 * np.exp(-((sto - b3) ** 2) / c3)\n",
    "        )\n",
    "        * 2.7\n",
    "    )\n",
    "\n",
    "    return D\n",
    "\n",
    "\n",
    "parameter_values.update({\"Positive particle diffusivity [m2.s-1]\": D_p})\n",
    "\n",
    "# add the diffusivity as a variable for easy plotting later\n",
    "sto = model.variables[\"Positive electrode stoichiometry\"]\n",
    "T = model.variables[\"Volume-averaged cell temperature [K]\"]\n",
    "model.variables[\"Positive particle diffusivity [m2.s-1]\"] = D_p(sto, T)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot what the diffusivity looks like as a function of stoichiometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sto = pybamm.linspace(0, 1, 100)\n",
    "_, ax = plt.subplots()\n",
    "ax = pybamm.plot(sto, D_p(sto, 298.15), ax=ax)\n",
    "ax.set_xlabel(\"Stoichiometry [-]\")\n",
    "ax.set_ylabel(\"Positive particle diffusivity [m2.s-1]\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we generate some synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters_for_data = parameter_values.copy()\n",
    "experiment = pybamm.Experiment(\n",
    "    [(\"Rest for 10s\", \"Discharge at C/10 for 5 minutes\", \"Rest for 30 minutes\")] * 10,\n",
    "    period=\"10 seconds\",\n",
    ")\n",
    "sim = iwp.Simulation(model, parameter_values=parameters_for_data, experiment=experiment)\n",
    "sim.solve(initial_soc=0.8)\n",
    "variables = [\"Time [s]\", \"Time [h]\", \"Current [A]\", \"Voltage [V]\"]\n",
    "data = pd.DataFrame(sim.solution.get_data_dict(variables))\n",
    "data = data.rename(columns={\"Cycle\": \"Cycle number\", \"Step\": \"Step number\"})\n",
    "t_h = data[\"Time [h]\"].values\n",
    "current = data[\"Current [A]\"].values\n",
    "data[\"Capacity [A.h]\"] = cumulative_trapezoid(current, t_h, initial=0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and plot the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.plot(\n",
    "    [\n",
    "        \"Current [A]\",\n",
    "        \"Voltage [V]\",\n",
    "        \"Positive electrode stoichiometry\",\n",
    "        \"Positive particle diffusivity [m2.s-1]\",\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we are generating synthetic data, we can also obtain the electrode stoichiometry. However, in practice we would need to calculate this from the capacity and voltage data, as well as the parameterization of the open-circuit voltage. The following block shows how to calculate the stoichiometry from the data.\n",
    "\n",
    "Note: this assumes we have already fitted the open-circuit potential and cell capacity, for example using the `MSMRHalfCell` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate initial stoichiometry from initial voltage based on the model OCP\n",
    "V_init = data[\"Voltage [V]\"].iloc[0]\n",
    "sto_calc = iwp.calculations.InitialStoichiometryFromVoltageHalfCell(\"positive\")\n",
    "params_for_pipeline = parameter_values.copy()\n",
    "params_for_pipeline.update(\n",
    "    {\"Initial voltage in positive electrode [V]\": V_init}, check_already_exists=False\n",
    ")\n",
    "sto_params = sto_calc.run(params_for_pipeline)\n",
    "sto_init = sto_params[\"Initial stoichiometry in positive electrode\"]\n",
    "# Calculate electrode capacity\n",
    "param = pybamm.LithiumIonParameters(options={\"working electrode\": \"positive\"})\n",
    "Q = params_for_pipeline.evaluate(param.p.prim.Q_init)\n",
    "# Add stoichiometry to the data\n",
    "data[\"Stoichiomety\"] = sto_init + data[\"Capacity [A.h]\"] / Q"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Fitting a scalar across all pulses\n",
    "\n",
    "The simplest assumption we can make is that the diffusivity takes a single scalar value across all pulses. We can fit this scalar value to the data. \n",
    "\n",
    "In the following we set up an objective for each cycle, passing in the initial concentration using the `custom_parameters` argument of the objective function. Any values provided in the `custom_parameter` dictionary are passed the simulation in that objective only. We create a dictionary of objectives, using the mid-point stoichiometry as the keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We fit a scalar for the diffusivity\n",
    "parameters = {\n",
    "    \"Positive particle diffusivity [m2.s-1]\": iwp.Parameter(\n",
    "        \"D_p\", initial_value=1e-14, bounds=(1e-15, 1e-13)\n",
    "    ),\n",
    "}\n",
    "\n",
    "# The model and solver are passed to the objective as options\n",
    "options = {\n",
    "    \"model\": model,\n",
    "}\n",
    "\n",
    "# We build a dictionary of objectives, indexed by the mid-point stoichiometry\n",
    "# We pass in the initial concentration in the positive electrode as a custom parameter\n",
    "# see the\n",
    "c_max = parameter_values[\"Maximum concentration in positive electrode [mol.m-3]\"]\n",
    "objectives = {}\n",
    "cycles = data[\"Cycle number\"].unique()\n",
    "for cycle in cycles:\n",
    "    data_cycle = data[data[\"Cycle number\"] == cycle]\n",
    "    # Use tht initial stoichiometry as to set the initial concentration\n",
    "    sto_init = data_cycle[\"Stoichiomety\"].iloc[0]\n",
    "    # Use midpoint stoichiometry as the stoichiometry for the cycle\n",
    "    sto_final = data_cycle[\"Stoichiomety\"].iloc[-1]\n",
    "    mid_point_sto = (sto_init + sto_final) / 2\n",
    "    # Create a new objective for each cycle\n",
    "    objective = iwp.objectives.Pulse(\n",
    "        data_cycle,\n",
    "        options=options,\n",
    "        custom_parameters={\n",
    "            \"Initial concentration in positive electrode [mol.m-3]\": c_max * sto_init\n",
    "        },\n",
    "    )\n",
    "    objectives[mid_point_sto] = objective\n",
    "\n",
    "# And create the data fit object\n",
    "scalar_data_fit = iwp.DataFit(objectives, parameters=parameters)\n",
    "\n",
    "# make sure we're not accidentally initializing with the correct values by passing\n",
    "# them in\n",
    "params_for_pipeline = {k: v for k, v in parameter_values.items() if k not in parameters}\n",
    "\n",
    "# Finally we run the fit\n",
    "scalar_results = scalar_data_fit.run(params_for_pipeline)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scalar_data_fit.plot_fit_results()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the fit is reasonable, but not perfect. We know the diffusivity is stoichiometry dependent, so let's look at some alternatives."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Fitting a scalar for each pulse\n",
    "\n",
    "We can instead assume that the diffusivity is a scalar but fit a value for each pulse in the data. This could then be used to create an interpolant or used as the basis for fitting a function to the calculated diffusivity.\n",
    "\n",
    "We can use the same objectives as before, but instead of setting up a `DataFit` object we use an `ArrayDataFit` object. This will fit the parameters separately for each objective in the dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array_data_fit = iwp.ArrayDataFit(objectives, parameters=parameters)\n",
    "array_results = array_data_fit.run(params_for_pipeline)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = array_data_fit.plot_fit_results()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This results in a much better fit for each individual pulse."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Calculating the diffusivity from each pulse\n",
    "\n",
    "We can calculate the diffusivity from each pulse using an analytical formula. The below pipeline calculation using equation 26 of [Wang, A. A., et al. \"Review of parameterisation and a novel database (LiionDB) for continuum Li-ion battery models.\" Progress in Energy 4.3 (2022)](https://iopscience.iop.org/article/10.1088/2516-1083/ac692c)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a dict to store sto vs diffusivity\n",
    "sto_D_data = {\"Stoichiometry\": [], \"Diffusivity [m2.s-1]\": []}\n",
    "\n",
    "# Loop over the cycles and calculate the diffusivity\n",
    "for cycle in data[\"Cycle number\"].unique():\n",
    "    data_cycle = data[data[\"Cycle number\"] == cycle]\n",
    "    # Use midpoint stoichiometry as the stoichiometry for the cycle\n",
    "    sto = (data_cycle[\"Stoichiomety\"].iloc[-1] + data_cycle[\"Stoichiomety\"].iloc[0]) / 2\n",
    "    # Calculate the diffusivity\n",
    "    calc_params_fit = iwp.calculations.DiffusivityFromPulse(\"positive\", data_cycle).run(\n",
    "        parameter_values\n",
    "    )\n",
    "    D = calc_params_fit[\"Positive particle diffusivity [m2.s-1]\"]\n",
    "    # Store the results\n",
    "    sto_D_data[\"Stoichiometry\"].append(sto)\n",
    "    sto_D_data[\"Diffusivity [m2.s-1]\"].append(D)\n",
    "\n",
    "sto_D_data = pd.DataFrame(sto_D_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then create an interpolant from our computed stoichiometry vs. diffusivity data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpolant = iwp.calculations.DiffusivityDataInterpolant(\"positive\", sto_D_data)\n",
    "interp_params = interpolant.run({})\n",
    "D_interp = interp_params[\"Positive particle diffusivity [m2.s-1]\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This approach often gives the correct \"shape\" for the stoichiometry dependent diffusivity, but the magnitude is often off. We can use this as a starting point for fitting a function to the data. In the following, we utilize the \"scale factor\" option which allows us to fit a scale factor that multiplies the calculated diffusivity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a new diffusivity function that scales the interpolated diffusivity\n",
    "scaled_interpolant = iwp.calculations.DiffusivityDataInterpolant(\n",
    "    \"positive\", sto_D_data, options={\"scale factor\": True}\n",
    ")\n",
    "\n",
    "\n",
    "# Set up the parameters for the fit\n",
    "parameters = {\n",
    "    \"Positive particle diffusivity scale factor\": iwp.Parameter(\n",
    "        \"D_scale\", bounds=(0.1, 10)\n",
    "    ),\n",
    "}\n",
    "\n",
    "\n",
    "# Create a pipeline that uses the interpolant and then fits the scale factor\n",
    "scaled_interp_data_fit = iwp.DataFit(objectives, parameters=parameters)\n",
    "pipeline = iwp.Pipeline(\n",
    "    {\n",
    "        \"known values\": iwp.direct_entries.DirectEntry(\n",
    "            params_for_pipeline, \"known values\"\n",
    "        ),\n",
    "        \"interpolant\": scaled_interpolant,\n",
    "        \"fit\": scaled_interp_data_fit,\n",
    "    }\n",
    ")\n",
    "scaled_interp_params_fit = pipeline.run()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scaled_interp_data_fit.plot_fit_results()"
   ]
  },
  {
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "6b12f1f625627d6c4a17ef696ddbbbd9bd4b12881121180de40e09e7956aa05c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}