{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Weighted cost functions\n",
    "\n",
    "In this notebook, we show different ways to weight the cost function when combining multiple objectives and output variables."
   ]
  },
  {
   "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",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate data\n",
    "\n",
    "First, we generate some synthetic data using a variant of the Robertson model. We add an activation energy term to the rate constants of the reaction so that we can run the simulation at different temperatures (assuming isothermal conditions). To keep things simple, we ignore units and assume the system is written in a suitable non-dimensional form."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Robertson(pybamm.BaseModel):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        x = pybamm.Variable(\"x\")\n",
    "        y = pybamm.Variable(\"y\")\n",
    "        z = pybamm.Variable(\"z\")\n",
    "\n",
    "        k_1 = pybamm.Parameter(\"k_1\")\n",
    "        E_1 = pybamm.Parameter(\"E_1\")\n",
    "        k_2 = pybamm.Parameter(\"k_2\")\n",
    "        E_2 = pybamm.Parameter(\"E_2\")\n",
    "        k_3 = pybamm.Parameter(\"k_3\")\n",
    "        E_3 = pybamm.Parameter(\"E_3\")\n",
    "        T = pybamm.Parameter(\"T\")\n",
    "\n",
    "        arr_1 = pybamm.exp(-E_1 / T)\n",
    "        arr_2 = pybamm.exp(-E_2 / T)\n",
    "        arr_3 = pybamm.exp(-E_3 / T)\n",
    "\n",
    "        self.rhs = {\n",
    "            x: -k_1 * arr_1 * x + k_2 * arr_2 * y * z,\n",
    "            y: k_1 * arr_1 * x - k_2 * arr_2 * y * z - k_3 * arr_3 * y**2,\n",
    "            z: k_3 * arr_3 * y**2,\n",
    "        }\n",
    "        self.initial_conditions = {x: 1, y: 0, z: 0}\n",
    "        self.variables = {\"x\": x, \"y\": y, \"z\": z}\n",
    "\n",
    "\n",
    "model = Robertson()\n",
    "true_inputs = {\n",
    "    \"k_1\": 5,\n",
    "    \"E_1\": 0.5,\n",
    "    \"k_2\": 3,\n",
    "    \"E_2\": 0.25,\n",
    "    \"k_3\": 1.5,\n",
    "    \"E_3\": 0.2,\n",
    "    \"T\": pybamm.InputParameter(\"T\"),\n",
    "}\n",
    "true_params = iwp.ParameterValues(true_inputs)\n",
    "true_params.process_model(model)\n",
    "solver = pybamm.IDAKLUSolver()\n",
    "t_data = np.linspace(0, 10, 1000)\n",
    "data_dict = {}\n",
    "for T in [1, 10]:\n",
    "    solution = solver.solve(model, t_eval=t_data, inputs={\"T\": T}, t_interp=t_data)\n",
    "    x_data = solution[\"x\"].data\n",
    "    y_data = solution[\"y\"].data\n",
    "    z_data = solution[\"z\"].data\n",
    "    data = pd.DataFrame({\"t\": t_data, \"x\": x_data, \"y\": y_data, \"z\": z_data})\n",
    "    data_dict[T] = data\n",
    "\n",
    "# Plot the data\n",
    "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\"]\n",
    "for i, (T, data) in enumerate(data_dict.items()):\n",
    "    plt.plot(data[\"t\"], data[\"x\"], \"-\", color=colors[i])\n",
    "    plt.plot(data[\"t\"], data[\"y\"], \"--\", color=colors[i])\n",
    "    plt.plot(data[\"t\"], data[\"z\"], \"-.\", color=colors[i])\n",
    "    if i == 0:\n",
    "        plt.plot([], [], \"-\", color=\"k\", label=\"x\")\n",
    "        plt.plot([], [], \"--\", color=\"k\", label=\"y\")\n",
    "        plt.plot([], [], \"-.\", color=\"k\", label=\"z\")\n",
    "    plt.plot([], [], \"o\", color=colors[i], label=f\"T = {T}\")\n",
    "plt.xlabel(\"t\")\n",
    "_ = plt.legend()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fit this model, we set up a custom `Objective` class which defines how the error function is calculated from the solution. Here we will compare the values of x, y, z between model and data. The cost function used to make the comparison can be specified later."
   ]
  },
  {
   "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\"], \"z\": data[\"z\"]}\n",
    "\n",
    "    def build(self, all_parameter_values):\n",
    "        data = self.data[\"data\"]\n",
    "        t_data = data[\"t\"].values\n",
    "\n",
    "        model = Robertson()\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, \"z\": sol[\"z\"].data}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a dictionary of objectives to fit the data at each temperature. When we do this, we pass in the known temperature as a `custom_parameter` to the objective. This is a parameter that is not fitted, but is passed to the model when the objective is evaluated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objectives = {}\n",
    "for T, data in data_dict.items():\n",
    "    objectives[T] = Objective(data, custom_parameters={\"T\": T})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We choose to use RMSE as our cost function. We can pass weights for both the objectives and variables to the `Cost` class. `objective_weights` should be a dictionary with the same keys as the objectives and the values are the weights for each objective. `variable_weights` should be a dictionary with the same keys as the variables and the values are the weights for each variable. If no weights are passed, the default is 1 for all objectives and variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective_weights = {1: 1, 10: 1}\n",
    "variable_weights = {\"x\": 0.1, \"y\": 0.1, \"z\": 0.8}\n",
    "cost = iwp.costs.RMSE(\n",
    "    objective_weights=objective_weights, variable_weights=variable_weights\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then specify which parameters need to be fitted, which parameters are known, define bounds, and pass this to the standard `DataFit` class to fit the unknown parameters. Let's assume we know $k_2$ and $E_2$ and fit the remaining parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_parameters = {\n",
    "    \"k_1\": iwp.Parameter(\"k_1\", initial_value=1, bounds=(0, 10)),\n",
    "    \"E_1\": iwp.Parameter(\"E_1\", initial_value=0.1, bounds=(0, 1)),\n",
    "    \"k_3\": iwp.Parameter(\"k_3\", initial_value=1, bounds=(0, 10)),\n",
    "    \"E_3\": iwp.Parameter(\"E_3\", initial_value=0.1, bounds=(0, 1)),\n",
    "}\n",
    "known_parameters = {\"k_2\": 3, \"E_2\": 0.5}\n",
    "\n",
    "optimizer = iwp.optimizers.ScipyMinimize()\n",
    "data_fit = iwp.DataFit(\n",
    "    objectives, parameters=fit_parameters, cost=cost, optimizer=optimizer\n",
    ")\n",
    "new_parameter_values = data_fit.run(known_parameters)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare the fitted and true values of the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in new_parameter_values.items():\n",
    "    print(f\"{key}: {value:.4f} (fit) {true_params[key]:.4f} (true)\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can 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(\n",
    "    axes[1:], [true_inputs[param.name] for param in data_fit.all_fit_parameters]\n",
    "):\n",
    "    ax.axhline(true_value, color=\"r\", linestyle=\"--\")"
   ]
  }
 ],
 "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
}