{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EIS fitting\n",
    "\n",
    "In this example, we show how to fit parameters using EIS data. This objective uses package [`pybammeis`](https://github.com/pybamm-team/pybamm-eis) to simulate the model response in the frequency domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pybammeis\n",
    "import pybamm\n",
    "import pandas as pd\n",
    "import ionworkspipeline as iwp\n",
    "import numpy as np"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate data\n",
    "\n",
    "We begin by generating synthetic EIS data for a Gr//SiOx half-cell. We will specify the reference exchange-current density and double-layer capacitance, and then simulate the impedance response for a range of frequencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model (DFN with capacitance)\n",
    "model = pybamm.lithium_ion.DFN(\n",
    "    options={\"working electrode\": \"positive\", \"surface form\": \"differential\"}\n",
    ")\n",
    "\n",
    "# Load parameters\n",
    "parameter_values = iwp.ParameterValues(\"OKane2022_graphite_SiOx_halfcell\")\n",
    "\n",
    "\n",
    "def j0(c_e, c_s_surf, c_s_max, T):\n",
    "    j0_ref = pybamm.Parameter(\n",
    "        \"Positive electrode reference exchange-current density [A.m-2]\"\n",
    "    )\n",
    "    c_e_init = pybamm.Parameter(\"Initial concentration in electrolyte [mol.m-3]\")\n",
    "\n",
    "    return (\n",
    "        j0_ref\n",
    "        * (c_e / c_e_init) ** 0.5\n",
    "        * (c_s_surf / c_s_max) ** 0.5\n",
    "        * (1 - c_s_surf / c_s_max) ** 0.5\n",
    "    )\n",
    "\n",
    "\n",
    "parameter_values.update(\n",
    "    {\n",
    "        \"Positive electrode reference exchange-current density [A.m-2]\": 2,\n",
    "        \"Positive electrode exchange-current density [A.m-2]\": j0,\n",
    "        \"Positive electrode double-layer capacity [F.m-2]\": 0.3,\n",
    "    },\n",
    "    check_already_exists=False,\n",
    ")\n",
    "\n",
    "# Create simulation\n",
    "eis_sim = pybammeis.EISSimulation(model, parameter_values=parameter_values)\n",
    "\n",
    "# Choose frequencies and calculate impedance\n",
    "frequencies = np.logspace(-4, 4, 30)\n",
    "sol = eis_sim.solve(frequencies)\n",
    "\n",
    "# Store data as arrays of real and complex impedance\n",
    "Z_Re = sol.real\n",
    "Z_Im = sol.imag\n",
    "data = pd.DataFrame(\n",
    "    {\"Frequency [Hz]\": frequencies, \"Z_Re [Ohm]\": Z_Re, \"Z_Im [Ohm]\": Z_Im}\n",
    ")\n",
    "\n",
    "# Generate a Nyquist plot\n",
    "_ = eis_sim.nyquist_plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit parameters\n",
    "\n",
    "We first set up the `EIS` objective. Then we pass in the data, model, and frequency range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = iwp.objectives.EIS(data, options={\"model\": model})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then specify which parameters need to be fitted (with initial guesses and bounds) and which parameters are known, and pass both to the standard `DataFit` class to fit the unknown parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"Positive electrode reference exchange-current density [A.m-2]\": iwp.Parameter(\n",
    "        \"j0_p_ref\", initial_value=1, bounds=(0.1, 10)\n",
    "    ),\n",
    "    \"Positive electrode double-layer capacity [F.m-2]\": iwp.Parameter(\n",
    "        \"C_dl_p\", initial_value=0.1, bounds=(0.01, 1)\n",
    "    ),\n",
    "}\n",
    "known_parameters = {k: v for k, v in parameter_values.items() if k not in parameters}\n",
    "\n",
    "data_fit = iwp.DataFit(objective, parameters=parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we run the fit and take a look at the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = data_fit.run(known_parameters)\n",
    "for name, p_fit in results.items():\n",
    "    print(f\"{name}: {p_fit} (fit) {parameter_values[name]} (true)\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot the results,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = data_fit.plot_fit_results()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and the trace of the cost and the parameter values to see how the fit progressed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = data_fit.plot_trace()"
   ]
  }
 ],
 "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
}