{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Fit example\n",
    "\n",
    "This notebook shows how the `DataFit` class works via examples. For design principles, see the [user_guide](../../user_guide/design_principles.md). For detailed documentation on the `DataFit` class, see the [API reference](../../api/data_fits/data_fits/index.rst).\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple non-battery example"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Ionworks battery parameter pipeline is built around PyBaMM models. We'll begin the demonstration of the `DataFit` class by creating a simple dynamical model that we will use to generate and fit synthetic 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",
    "\n",
    "\n",
    "class LotkaVolterra(pybamm.BaseModel):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        x = pybamm.Variable(\"x\")\n",
    "        y = pybamm.Variable(\"y\")\n",
    "\n",
    "        a = pybamm.Parameter(\"a\")\n",
    "        b = pybamm.Parameter(\"b\")\n",
    "        c = pybamm.Parameter(\"c\")\n",
    "        d = pybamm.Parameter(\"d\")\n",
    "\n",
    "        self.rhs = {x: a * x - b * x * y, y: -c * y + d * x * y}\n",
    "        self.initial_conditions = {x: 2, y: 1}\n",
    "        self.variables = {\"x\": x, \"y\": y}\n",
    "\n",
    "\n",
    "model = LotkaVolterra()\n",
    "true_inputs = {\"a\": 1.5, \"b\": 0.5, \"c\": 3.0, \"d\": 0.3}\n",
    "true_params = iwp.ParameterValues(true_inputs)\n",
    "true_params.process_model(model)\n",
    "solver = pybamm.ScipySolver()\n",
    "t_data = np.linspace(0, 10, 1000)\n",
    "solution = solver.solve(model, t_data)\n",
    "\n",
    "# extract data and add noise\n",
    "x_data = solution[\"x\"].data + 0.1 * np.random.rand(t_data.size)\n",
    "y_data = solution[\"y\"].data + 0.1 * np.random.rand(t_data.size)\n",
    "\n",
    "data = pd.DataFrame({\"t\": t_data, \"x\": x_data, \"y\": y_data})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fit this model, we set up a custom `Objective` class that defines how the error function is calculated from the solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Objective(iwp.objectives.Objective):\n",
    "    def process_data(self):\n",
    "        data = self.data[\"data\"]\n",
    "        self._processed_data = {\"x\": data[\"x\"], \"y\": data[\"y\"]}\n",
    "\n",
    "    def build(self, all_parameter_values):\n",
    "        data = self.data[\"data\"]\n",
    "        t_data = data[\"t\"].values\n",
    "\n",
    "        model = LotkaVolterra()\n",
    "        all_parameter_values.process_model(model)\n",
    "        self.simulation = iwp.Simulation(model, t_eval=t_data, t_interp=t_data)\n",
    "\n",
    "    def run(self, inputs, full_output=False):\n",
    "        sim = self.simulation\n",
    "        sol = sim.solve(inputs=inputs)\n",
    "        return {\"x\": sol[\"x\"].data, \"y\": sol[\"y\"].data}\n",
    "\n",
    "\n",
    "objective = Objective(data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then specify the parameters that need to be fitted, identify any known parameters (none in this case), define bounds, and pass these specifications to the standard DataFit class to fit the unknown parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_parameters = {\n",
    "    \"a\": iwp.Parameter(\"a\", bounds=(0, np.inf)),\n",
    "    \"b\": iwp.Parameter(\"b\", bounds=(0, np.inf)),\n",
    "    \"c\": iwp.Parameter(\"c\", bounds=(0, np.inf)),\n",
    "    \"d\": iwp.Parameter(\"d\", bounds=(0, np.inf)),\n",
    "}\n",
    "known_parameters = {}\n",
    "\n",
    "optimizer = iwp.optimizers.ScipyLeastSquares(verbose=2)\n",
    "data_fit = iwp.DataFit(objective, parameters=fit_parameters, optimizer=optimizer)\n",
    "new_parameter_values = data_fit.run(known_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then plot the trace of the cost and the parameter values to see how the fit progressed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = data_fit.plot_trace()\n",
    "for ax, true_value in zip(axes[1:], true_inputs.values()):\n",
    "    ax.axhline(true_value, color=\"r\", linestyle=\"--\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can reuse the same `Objective` and `DataFit` class but change which parameters are known and which are to be fitted. This flexibility is useful for battery models where we want to fit only a subset of the parameters to experimental data while keeping other parameters fixed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_parameters = {\n",
    "    \"a\": iwp.Parameter(\"a\", bounds=(0, np.inf)),\n",
    "    \"b\": iwp.Parameter(\"b\", bounds=(0, np.inf)),\n",
    "}\n",
    "known_parameters = {\"c\": 3, \"d\": 0.3}\n",
    "\n",
    "data_fit = iwp.DataFit(objective, parameters=fit_parameters, optimizer=optimizer)\n",
    "new_parameter_values = data_fit.run(known_parameters)\n",
    "\n",
    "for k, v in new_parameter_values.items():\n",
    "    true = true_inputs[k]\n",
    "    print(f\"Parameter {k}: true: {true:.3f}, fitted: {v:.3f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Battery example: constant-current experiment"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now repeat the same procedure for a battery model. We start by creating a model and solving it to generate synthetic data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = pybamm.lithium_ion.SPM()\n",
    "parameter_values = iwp.ParameterValues(\"Chen2020\")\n",
    "sim = iwp.Simulation(model, parameter_values=parameter_values)\n",
    "t = np.linspace(0, 3600, 1000)\n",
    "sim.solve([0, 3600], t_interp=t)\n",
    "data = pd.DataFrame(\n",
    "    {x: sim.solution[x].entries for x in [\"Time [s]\", \"Current [A]\", \"Voltage [V]\"]}\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we set up which parameters are to be fitted, and call the pre-defined `CurrentDriven` objective function, which solves the model with the current from the data and calculates the error between voltage from the model and the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example we just fit the diffusivity in the positive electrode\n",
    "parameters = {\n",
    "    \"Positive particle diffusivity [m2.s-1]\": iwp.Parameter(\"D_s\", initial_value=1e-15),\n",
    "}\n",
    "\n",
    "objective = iwp.objectives.CurrentDriven(data, options={\"model\": model})\n",
    "current_driven = iwp.DataFit(objective, 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",
    "results = current_driven.run(params_for_pipeline)\n",
    "\n",
    "print(\n",
    "    f\"True parameter value: {parameter_values['Positive particle diffusivity [m2.s-1]']:.3e}\"\n",
    ")\n",
    "print(\n",
    "    f\"Fitted parameter value: {results['Positive particle diffusivity [m2.s-1]']:.3e}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `results` variable is a `Result` class which contains information about the optimization. The optimal parameter values are accessible by indexing the solution, like\n",
    "\n",
    "```python\n",
    "results['Positive particle diffusivity [m2.s-1]']\n",
    "```\n",
    "\n",
    "or through its properties,\n",
    "\n",
    "- `parameter_values`: The final values of the optimized parameters.\n",
    "- `optimizer_result`: The result object returned by the optimizer.\n",
    "- `callbacks`: The callbacks used during optimization.\n",
    "- `children`: Results from multistarted parameter estimation (if multistart was used)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.parameter_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot the fit results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_driven.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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}