{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combining Objectives example\n",
    "\n",
    "This notebook shows how to construct and combine custom `Objective` classes to fit multiple objectives at once (e.g. simultaneously fit two data sets). For design principles, see the [user_guide](../../user_guide/design_principles.md). For detailed documentation on the `Objective` class, see the [API reference](../../api/data_fits/objectives/index.rst)."
   ]
  },
  {
   "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple pulse example\n",
    "\n",
    "For this example, we set up a single particle model with a temperature dependent diffusivity. The goal is to fit the activation energy of the diffusivity using data generated from simple pulse experiments conducted at different temperatures.\n",
    "\n",
    "We first set up the model and parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = pybamm.lithium_ion.SPM()\n",
    "parameter_values = iwp.ParameterValues(\"Chen2020\")\n",
    "\n",
    "\n",
    "def D_n(c, T):\n",
    "    D_ref = pybamm.Parameter(\"Negative electrode reference diffusivity [m2.s-1]\")\n",
    "    E_D_s = pybamm.Parameter(\n",
    "        \"Negative particle diffusivity activation energy [J.mol-1]\"\n",
    "    )\n",
    "    arrhenius = np.exp(-E_D_s / (pybamm.constants.R * T)) * np.exp(\n",
    "        E_D_s / (pybamm.constants.R * 296)\n",
    "    )\n",
    "    return D_ref * arrhenius\n",
    "\n",
    "\n",
    "parameter_values.update(\n",
    "    {\n",
    "        \"Negative particle diffusivity [m2.s-1]\": D_n,\n",
    "        \"Negative electrode reference diffusivity [m2.s-1]\": 3.3e-14,\n",
    "        \"Negative particle diffusivity activation energy [J.mol-1]\": 2.5e4,\n",
    "    },\n",
    "    check_already_exists=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we generate some synthetic data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dict = {}\n",
    "sols = []\n",
    "for T in [287.15, 298.15, 318.15]:\n",
    "    parameters_for_data = parameter_values.copy()\n",
    "    experiment = pybamm.Experiment(\n",
    "        [\"Discharge at C/3 for 10 minutes\", \"Rest for 15 minutes\"],\n",
    "        period=\"10 seconds\",\n",
    "        temperature=T,\n",
    "    )\n",
    "    sim = pybamm.Simulation(\n",
    "        model, parameter_values=parameters_for_data, experiment=experiment\n",
    "    )\n",
    "    sim.solve()\n",
    "    sols.append(sim.solution)\n",
    "    data_dict[T] = pd.DataFrame(\n",
    "        {x: sim.solution[x].entries for x in [\"Time [s]\", \"Current [A]\", \"Voltage [V]\"]}\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look at the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pybamm.dynamic_plot(\n",
    "    sols, [\"Current [A]\", \"Voltage [V]\"], labels=[f\"{T} K\" for T in data_dict]\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now 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. We can set up an objective for the data collected at each temperature, passing in the temperature 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example we fit both the reference diffusivity and activation energy in\n",
    "# the negative electrode\n",
    "parameters = {\n",
    "    \"Negative electrode reference diffusivity [m2.s-1]\": iwp.Parameter(\n",
    "        \"D_ref\", initial_value=1e-13, bounds=(1e-15, 1e-12)\n",
    "    ),\n",
    "    \"Negative particle diffusivity activation energy [J.mol-1]\": iwp.Parameter(\n",
    "        \"E_D_s\", initial_value=1e4, bounds=(1e3, 1e5)\n",
    "    ),\n",
    "}\n",
    "\n",
    "# We loop over the data at different temperatures, and create a separate objective\n",
    "# for each temperature, storing them in a dictionary\n",
    "objectives = {}\n",
    "for T, data in data_dict.items():\n",
    "    objectives[f\"{T} K\"] = iwp.objectives.CurrentDriven(\n",
    "        data, options={\"model\": model}, custom_parameters={\"Ambient temperature [K]\": T}\n",
    "    )\n",
    "\n",
    "# We then create a data fit object, passing in the objectives and the parameters. This\n",
    "# will fit the parameters to all data simultaneously\n",
    "current_driven = 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",
    "results = current_driven.run(params_for_pipeline)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for parameter in parameters:\n",
    "    print(parameter)\n",
    "    print(f\"True parameter value: {parameter_values[parameter]:.3e}\")\n",
    "    print(f\"Fitted parameter value: {results[parameter]:.3e}\")"
   ]
  },
  {
   "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
}