{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pipeline serialization\n",
    "\n",
    "In this example, we will show how to serialize and deserialize pipeline elements for saving or use with the API. The datafit is the same as is constructed in `particle_swarm_optimiser.ipynb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ionworkspipeline as iwp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pybamm\n",
    "\n",
    "# Load the model and parameter values\n",
    "model = pybamm.lithium_ion.SPM()\n",
    "parameter_values = iwp.ParameterValues(\"Chen2020\")\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",
    "# Generate the data\n",
    "t = np.arange(0, 900, 3)\n",
    "sim = iwp.Simulation(\n",
    "    model, parameter_values=parameter_values, t_eval=[0, 900], 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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define the parameters to fit, along with initial guesses and bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"Negative electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Negative electrode active material volume fraction\",\n",
    "        initial_value=0.5,\n",
    "        bounds=(0, 1),\n",
    "    ),\n",
    "    \"Positive electrode active material volume fraction\": iwp.Parameter(\n",
    "        \"Positive electrode active material volume fraction\",\n",
    "        initial_value=0.5,\n",
    "        bounds=(0, 1),\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set up the objective function, which in this case is the current driven objective. This takes the time vs current data and runs the model, comparing the model's predictions for the voltage with the experimental data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = iwp.objectives.CurrentDriven(data, options={\"model\": model})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we set up the `DataFit` object, which takes the objective function, the parameters to fit, and the optimizer as inputs. We also convert the parameters to a direct entry. Finally, we construct a pipeline to demonstrate the serialization and deserialization."
   ]
  },
  {
   "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",
    "params_direct_entry = iwp.direct_entries.DirectEntry(\n",
    "    params_for_pipeline, source=\"Chen2020\"\n",
    ")\n",
    "data_fit = iwp.DataFit(\n",
    "    objective,\n",
    "    parameters=parameters,\n",
    "    options={\"maxiters\": 100, \"seed\": 0},\n",
    ")\n",
    "\n",
    "pipeline = iwp.Pipeline(\n",
    "    {\"starting parameters\": params_direct_entry, \"data fit\": data_fit}\n",
    ")"
   ]
  },
  {
   "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": [
    "results = pipeline.run()\n",
    "print(\n",
    "    f\"Negative electrode active material volume fraction: {results['Negative electrode active material volume fraction']}\"\n",
    ")\n",
    "print(\n",
    "    f\"Positive electrode active material volume fraction: {results['Positive electrode active material volume fraction']}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we convert the pipeline to a JSON configuration file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_config = pipeline.to_config()\n",
    "# To save the JSON config to a file and view:\n",
    "# import json\n",
    "# with open(\"pipeline_config.json\", \"w\") as f:\n",
    "#     json.dump(json_config, f, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and convert it back into a pipeline element. Then we can run it, and check whether we get the same results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_pipeline = iwp.parsers.parse_pipeline(json_config)\n",
    "results = new_pipeline.run()\n",
    "print(\n",
    "    f\"Negative electrode active material volume fraction: {results['Negative electrode active material volume fraction']}\"\n",
    ")\n",
    "print(\n",
    "    f\"Positive electrode active material volume fraction: {results['Positive electrode active material volume fraction']}\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}