{
 "cells": [
  {
   "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 ionworkspipeline as iwp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pybammeis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we generate timeseries data for a `CurrentDriven` objective using the `SPM` with the differential surface form option enabled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model and parameter values\n",
    "model = pybamm.lithium_ion.SPM(options={\"surface form\": \"differential\"})\n",
    "parameter_values = iwp.ParameterValues(\"Chen2020\")\n",
    "\n",
    "parameter_values.update(\n",
    "    {\n",
    "        \"Negative electrode active material volume fraction\": 0.75,\n",
    "        \"Positive electrode active material volume fraction\": 0.665,\n",
    "    }\n",
    ")\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",
    "# Generate the data\n",
    "t = np.arange(0, 3000, 3)\n",
    "sim = iwp.Simulation(\n",
    "    model, parameter_values=parameter_values, t_eval=[t[0], t[-1]], t_interp=t\n",
    ")\n",
    "sol = sim.solve()\n",
    "data = pd.DataFrame(\n",
    "    {x: sim.solution[x](t) for x in [\"Time [s]\", \"Current [A]\", \"Voltage [V]\"]}\n",
    ")\n",
    "# add noise to the data\n",
    "sigma = 0.001\n",
    "\n",
    "# Set seed for reproducibility\n",
    "np.random.seed(0)\n",
    "data[\"Voltage [V]\"] += np.random.normal(0, sigma, len(data))\n",
    "\n",
    "objective_current_driven = iwp.objectives.CurrentDriven(data, options={\"model\": model})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we use the same `model` and `parameter_values` to create synthetic EIS data and the `EIS` objective."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "objective_eis = iwp.objectives.EIS(data, options={\"model\": model})\n",
    "\n",
    "# Generate a Nyquist plot\n",
    "_ = eis_sim.nyquist_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we define the fitting parameters that will be optimized using both the `CurrentDriven` and `EIS` data. We perform a `Log10` transform on the cathode exchange-current density and double-layer capacity because these parameters span several orders of magnitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"Negative electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Anode active material frac.\",\n",
    "        bounds=(0.3, 0.8),\n",
    "    ),\n",
    "    \"Positive electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Cathode active material frac.\",\n",
    "        bounds=(0.3, 0.8),\n",
    "    ),\n",
    "    \"Positive electrode reference exchange-current density [A.m-2]\": iwp.transforms.Log10(\n",
    "        iwp.Parameter(\n",
    "            \"j0_p_ref\",\n",
    "            bounds=(0.1, 10),\n",
    "        )\n",
    "    ),\n",
    "    \"Positive electrode double-layer capacity [F.m-2]\": iwp.transforms.Log10(\n",
    "        iwp.Parameter(\n",
    "            \"C_dl_p\",\n",
    "            bounds=(0.01, 1),\n",
    "        )\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we combine the two objectives to optimize them simultaneously. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objectives = {\n",
    "    \"CurrentDriven\": objective_current_driven,\n",
    "    \"EIS\": objective_eis,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All `Objective`s have a default optimizer associated with them that is tuned for the particular usecase. A specific optimizer can always be given to the `DataFit` class, but if none is provided, the optimizer will default to the one given by the first objective.\n",
    "\n",
    "The default optimizers for the `EIS` and `CurrentDriven` objectives are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, objective in objectives.items():\n",
    "    print(f\"{name}: {objective.default_optimizer.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we create the `DataFit` class combining both objectives. The `DataFit` uses the `SciPy Minimize Optimizer` provided by `objective_current_driven`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fit = iwp.DataFit(objectives, parameters=parameters)\n",
    "\n",
    "print(f\"DataFit optimizer: {data_fit.optimizer.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we run the data fit, passing the parameters that are not being fit as a dictionary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_for_pipeline = {k: v for k, v in parameter_values.items() if k not in parameters}\n",
    "results = data_fit.run(params_for_pipeline)\n",
    "for k, v in results.items():\n",
    "    print(f\"{k}: {parameter_values[k]:.3e} (true) {v:.3e} (fit)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we plot the results and the trace of the cost function and parameter values for both the `CurrentDriven` and `EIS` objectives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = data_fit.plot_fit_results()\n",
    "fig, axes = data_fit.plot_trace()\n",
    "\n",
    "sorted_keys = np.asarray(list(parameters.keys()))[\n",
    "    np.argsort([p.name for p in parameters.values()])\n",
    "]\n",
    "\n",
    "for ax, k in zip(axes[1:], sorted_keys):\n",
    "    v = parameter_values[k]\n",
    "    if data_fit.parameters_original[k].is_transformed:\n",
    "        v = data_fit.parameters_original[k].transform(v)\n",
    "    ax.axhline(v, color=\"red\", linestyle=\"--\", label=\"True\")\n",
    "    ax.legend()"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}