{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell",
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# WARNING: advised to install a specific version, e.g. tensorwaves==0.1.2\n",
    "%pip install -q tensorwaves[doc,jax,viz] IPython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "import os\n",
    "\n",
    "from IPython.display import display  # noqa: F401\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Perform fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As explained in the {doc}`previous step <step2>`, a {class}`.Model` object and its lambdified {class}`.Function` form behave just like a mathematical function that takes a set of data points as an argument and returns a list of intensities (real numbers). At this stage, we want to optimize the parameters of the intensity model, so that it matches the distribution of our data sample. This is what we call 'performing a fit'.\n",
    "\n",
    "First, load the relevant data from the previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "import qrules as q\n",
    "\n",
    "from tensorwaves.model import SympyModel\n",
    "\n",
    "result = q.io.load(\"transitions.json\")\n",
    "with open(\"helicity_model.pickle\", \"rb\") as stream:\n",
    "    model = pickle.load(stream)\n",
    "with open(\"data_set.pickle\", \"rb\") as stream:\n",
    "    data_set = pickle.load(stream)\n",
    "with open(\"phsp_set.pickle\", \"rb\") as stream:\n",
    "    phsp_set = pickle.load(stream)\n",
    "sympy_model = SympyModel(\n",
    "    expression=model.expression,\n",
    "    parameters=model.parameter_defaults,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Define estimator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform a fit, you need to define an {class}`.Estimator`. An estimator is a measure for the discrepancy between the intensity model and the data distribution to which you fit. In PWA, we usually use an *unbinned negative log likelihood estimator*.\n",
    "\n",
    "Generally, the intensity model is not normalized, but a log likelihood estimator requires a normalized function. This is where the {ref}`phase space dataset <usage/step2:2.1 Generate phase space sample>` comes into play again: the intensity is evaluated separately with the phase space dataset so that the output of the {class}`.Function` can be normalized with it. The phase space sample is therefore a required argument!\n",
    "\n",
    "```{margin}\n",
    "If you want to correct for the efficiency of the detector, you should insert a *detector-reconstructed* phase space sample here.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.estimator import UnbinnedNLL\n",
    "\n",
    "estimator = UnbinnedNLL(\n",
    "    sympy_model,\n",
    "    data_set,\n",
    "    phsp_set,\n",
    "    backend=\"jax\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the {class}`.UnbinnedNLL` can be expressed with different backends (it creates a {class}`.LambdifiedFunction` internally from the template {class}`.Model`). Here, we use {func}`jax <jax.jit>`, which is turns out to be the fastest backend for this model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Optimize intensity model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Specify fit parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting the fit itself is quite simple: just create an optimizer instance of your choice, here [Minuit2](https://root.cern.ch/doc/master/Minuit2Page.html), and call its {meth}`~.Minuit2.optimize` method to start the fitting process. The {meth}`~.Minuit2.optimize` method requires a mapping of parameter names to their initial values. **Only the parameters listed in the mapping are optimized.**\n",
    "\n",
    "Let's first select a few of the parameters that we saw in {ref}`Step 3.1 <usage/step3:3.1 Define estimator>` and feed them to the optimizer to run the fit. Notice that we modify the parameters slightly to make the fit more interesting (we are running fitting to a data sample that was generated with this very same amplitude model after all)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_parameters = {\n",
    "    \"C[J/\\\\psi(1S) \\\\to f_{0}(1500)_{0} \\\\gamma_{+1}; f_{0}(1500) \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0}]\": 1.0\n",
    "    + 0.0j,\n",
    "    \"Gamma_f(0)(980)\": 0.15,\n",
    "    \"Gamma_f(0)(1500)\": 0.2,\n",
    "    \"m_f(0)(1710)\": 1.78,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that a {class}`.Function` object computes the intensity for a certain dataset. This can be seen now nicely when we use these intensities as weights on the phase space sample and plot it together with the original dataset. Here, we look at the invariant mass distribution projection of the final states `1` and `2`, which, {ref}`as we saw before <usage/step2:2.3 Visualize kinematic variables>`, are the final state particles $\\pi^0,\\pi^0$.\n",
    "\n",
    "Don't forget to use {meth}`~.Function.update_parameters` first!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib import cm\n",
    "\n",
    "reaction_info = model.adapter.reaction_info\n",
    "intermediate_states = sorted(\n",
    "    (\n",
    "        p\n",
    "        for p in model.particles\n",
    "        if p not in reaction_info.final_state.values()\n",
    "        and p not in reaction_info.initial_state.values()\n",
    "    ),\n",
    "    key=lambda p: p.mass,\n",
    ")\n",
    "\n",
    "evenly_spaced_interval = np.linspace(0, 1, len(intermediate_states))\n",
    "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n",
    "\n",
    "\n",
    "def indicate_masses():\n",
    "    plt.xlabel(\"$m$ [GeV]\")\n",
    "    for i, p in enumerate(intermediate_states):\n",
    "        plt.gca().axvline(\n",
    "            x=p.mass, linestyle=\"dotted\", label=p.name, color=colors[i]\n",
    "        )\n",
    "\n",
    "\n",
    "def compare_model(\n",
    "    variable_name,\n",
    "    data_set,\n",
    "    phsp_set,\n",
    "    intensity_model,\n",
    "    bins=100,\n",
    "):\n",
    "    data = data_set[variable_name]\n",
    "    phsp = phsp_set[variable_name]\n",
    "    intensities = intensity_model(phsp_set)\n",
    "    axis = plt.gca()\n",
    "    axis.hist(data, bins=bins, alpha=0.5, label=\"data\", density=True)\n",
    "    axis.hist(\n",
    "        phsp,\n",
    "        weights=intensities,\n",
    "        bins=bins,\n",
    "        histtype=\"step\",\n",
    "        color=\"red\",\n",
    "        label=\"initial fit model\",\n",
    "        density=True,\n",
    "    )\n",
    "    indicate_masses()\n",
    "    plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.model import LambdifiedFunction\n",
    "\n",
    "intensity = LambdifiedFunction(sympy_model, backend=\"numpy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity.update_parameters(initial_parameters)\n",
    "compare_model(\"m_12\", data_set, phsp_set, intensity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Custom callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {class}`.Minuit2` class allows one to insert certain {mod}`~tensorwaves.optimizer.callbacks`. Callbacks are used to insert behavior into the {meth}`.Optimizer.optimize` method. The {mod}`~tensorwaves.optimizer.callbacks` module provides several handy callback classes, but it's also possible to define custom callbacks into the {class}`.Optimizer` by defining a custom {class}`.Callback` class. Here's one that live updates a plot of the latest fit model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "from tensorwaves.optimizer.callbacks import Callback\n",
    "\n",
    "variable_to_plot = \"m_12\"\n",
    "is_doc = \"READTHEDOCS\" in os.environ or \"EXECUTE_NB\" in os.environ\n",
    "\n",
    "\n",
    "class PyplotCallback(Callback):\n",
    "    def __init__(self, step_size=10):\n",
    "        self.__step_size = step_size\n",
    "        self.__fig = None\n",
    "        self.__ax = None\n",
    "        self.__latest_parameters = {}\n",
    "\n",
    "    def on_optimize_start(self, logs):\n",
    "        self.__fig, self.__ax = plt.subplots(1, figsize=(8, 5))\n",
    "\n",
    "    def on_optimize_end(self, logs):\n",
    "        self.update_plot(nbins=200)\n",
    "        self.__ax = None\n",
    "        self.__fig = None\n",
    "\n",
    "    def on_iteration_end(self, iteration, logs=None):\n",
    "        pass\n",
    "\n",
    "    def on_function_call_end(self, function_call, logs):\n",
    "        self.__latest_parameters = logs[\"parameters\"]\n",
    "        if is_doc:\n",
    "            return\n",
    "        if function_call % self.__step_size != 0:\n",
    "            return\n",
    "        if function_call > 75:\n",
    "            return\n",
    "        self.update_plot(nbins=80)\n",
    "        clear_output(wait=True)\n",
    "        display(plt.gcf())\n",
    "\n",
    "    def update_plot(self, nbins: int):\n",
    "        if self.__fig is None or self.__ax is None:\n",
    "            self.__fig, self.__ax = plt.subplots(1, figsize=(8, 5))\n",
    "        data = data_set[variable_to_plot]\n",
    "        phsp = phsp_set[variable_to_plot]\n",
    "        intensity.update_parameters(self.__latest_parameters)\n",
    "        intensities = intensity(phsp_set)\n",
    "        self.__ax.cla()\n",
    "        self.__ax.hist(data, bins=nbins, alpha=0.5, label=\"data\", density=True)\n",
    "        self.__ax.hist(\n",
    "            phsp,\n",
    "            weights=intensities,\n",
    "            bins=nbins,\n",
    "            histtype=\"step\",\n",
    "            color=\"red\",\n",
    "            label=\"fit model\",\n",
    "            density=True,\n",
    "        )\n",
    "        self.__ax.set_xlim((0.25, 2.5))\n",
    "        self.__ax.set_ylim((0, 1.9))\n",
    "        indicate_masses()\n",
    "        plt.gcf().legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the fit example below, we want to use several callbacks together. To do so, we 'stack' them together with {class}`.CallbackList`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.optimizer.callbacks import (\n",
    "    CallbackList,\n",
    "    CSVSummary,\n",
    "    TFSummary,\n",
    "    YAMLSummary,\n",
    ")\n",
    "\n",
    "callbacks = CallbackList(\n",
    "    [\n",
    "        CSVSummary(\"fit_traceback_minuit.csv\"),\n",
    "        PyplotCallback(),  # comment this line for raw fit performance\n",
    "        TFSummary(),\n",
    "        YAMLSummary(\"current_fit_result.yaml\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Perform fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we create an {class}`.Optimizer` to {meth}`~.Minuit2.optimize` the model (which is embedded in the {class}`.Estimator`). Here, we choose the {class}`.Minuit2` optimizer, which is the most common optimizer in high-energy physics. Since we constructed the {class}`.UnbinnedNLL` with {obj}`jax <jax.jit>`, can an optimize the model with analytic gradient over the {class}`.LambdifiedFunction`:\n",
    "\n",
    "```{margin}\n",
    "The computation time depends on the complexity of the model, number of data events, the size of the phase space sample, and the number of free parameters. This model is rather small and has but a few free parameters, so the optimization shouldn't take more than a **few seconds**.\n",
    "\n",
    "The custom callback example slows down the optimization, but you can remove it to get the raw fit performance.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.optimizer.minuit import Minuit2\n",
    "\n",
    "minuit2 = Minuit2(\n",
    "    callback=callbacks,\n",
    "    use_analytic_gradient=False,\n",
    ")\n",
    "fit_result = minuit2.optimize(estimator, initial_parameters)\n",
    "fit_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, the values of the optimized parameters in the result are again comparable to the original values that are contained in the model ({attr}`.SympyModel.parameters`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "optimized_parameters = fit_result.parameter_values\n",
    "for p in optimized_parameters:\n",
    "    print(p)\n",
    "    print(f\"  initial:   {initial_parameters[p]:.3}\")\n",
    "    print(f\"  optimized: {optimized_parameters[p]:.3}\")\n",
    "    print(f\"  original:  {sympy_model.parameters[p]:.3}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{ref}`usage/step3:SciPy optimizer`\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 3.3 Export and import"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In {ref}`usage/step3:3.2 Optimize intensity model`, we initialized {obj}`.Minuit2` with some callbacks that are {class}`.Loadable`. Such callback classes offer the possibility to {meth}`.Loadable.load_latest_parameters`, so you can pick up the optimize process in case it crashes or if you pause it. Loading the latest parameters goes as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latest_parameters = YAMLSummary.load_latest_parameters(\n",
    "    \"current_fit_result.yaml\"\n",
    ")\n",
    "latest_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To restart the fit with the latest parameters, simply rerun as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minuit2 = Minuit2()\n",
    "minuit2.optimize(estimator, latest_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lo and behold: the parameters were already optimized, so the fit converged faster!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.4 Visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot optimized model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the same method as above, we renew the parameters of the {class}`.Function` and plot it again over the data distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity.update_parameters(latest_parameters)\n",
    "compare_model(\"m_12\", data_set, phsp_set, intensity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Intensity components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the {class}`~ampform.helicity.HelicityModel` contains {attr}`~ampform.helicity.HelicityModel.components` attribute. This is {obj}`dict` of component names to {class}`sympy.Expr <sympy.core.expr.Expr>`s helps us to identify amplitudes and intensities in the total amplitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(model.components)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.components[\n",
    "    R\"A[J/\\psi(1S)_{+1} \\to f_{0}(1370)_{0} \\gamma_{+1}; f_{0}(1370)_{0} \\to\"\n",
    "    R\" \\pi^{0}_{0} \\pi^{0}_{0}]\"\n",
    "].subs(model.parameter_defaults).doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like in {ref}`usage/step2:2.2 Generate intensity-based sample`, these _intensity components_ can each be expressed in a computational backend. This can be done with the helper function {func}`.create_intensity_component`. After that we, update the parameters with the optimized parameter values we found in {ref}`usage/step3:Perform fit`. The two components in the example below should be the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.physics import create_intensity_component\n",
    "\n",
    "from_amplitudes = create_intensity_component(\n",
    "    model,\n",
    "    components=[\n",
    "        R\"A[J/\\psi(1S)_{+1} \\to f_{0}(500)_{0} \\gamma_{+1}; f_{0}(500)_{0} \\to\"\n",
    "        R\" \\pi^{0}_{0} \\pi^{0}_{0}]\",\n",
    "        R\"A[J/\\psi(1S)_{+1} \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980)_{0} \\to\"\n",
    "        R\" \\pi^{0}_{0} \\pi^{0}_{0}]\",\n",
    "        R\"A[J/\\psi(1S)_{+1} \\to f_{0}(1370)_{0} \\gamma_{+1}; f_{0}(1370)_{0}\"\n",
    "        R\" \\to \\pi^{0}_{0} \\pi^{0}_{0}]\",\n",
    "        R\"A[J/\\psi(1S)_{+1} \\to f_{0}(1500)_{0} \\gamma_{+1}; f_{0}(1500)_{0}\"\n",
    "        R\" \\to \\pi^{0}_{0} \\pi^{0}_{0}]\",\n",
    "        R\"A[J/\\psi(1S)_{+1} \\to f_{0}(1710)_{0} \\gamma_{+1}; f_{0}(1710)_{0}\"\n",
    "        R\" \\to \\pi^{0}_{0} \\pi^{0}_{0}]\",\n",
    "    ],\n",
    "    backend=\"numpy\",\n",
    ")\n",
    "from_amplitudes.update_parameters(latest_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from_intensity = create_intensity_component(\n",
    "    model,\n",
    "    components=R\"I[J/\\psi(1S)_{+1} \\to \\gamma_{+1} \\pi^{0}_{0} \\pi^{0}_{0}]\",\n",
    "    backend=\"numpy\",\n",
    ")\n",
    "from_intensity.update_parameters(latest_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "difference = np.average(from_amplitudes(phsp_set) - from_intensity(phsp_set))\n",
    "assert np.round(difference, decimals=15) == 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a {class}`.LambdifiedFunction` that can be plotted just like in {ref}`usage/step3:Plot optimized model`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(8, 5))\n",
    "bins = 150\n",
    "ax.hist(\n",
    "    phsp_set[\"m_12\"],\n",
    "    weights=intensity(phsp_set),\n",
    "    bins=bins,\n",
    "    alpha=0.2,\n",
    "    label=\"full intensity\",\n",
    ")\n",
    "ax.hist(\n",
    "    phsp_set[\"m_12\"],\n",
    "    weights=from_intensity(phsp_set),\n",
    "    bins=bins,\n",
    "    histtype=\"step\",\n",
    "    label=R\"$J/\\psi(1S)_{-1} \\to \\gamma_{-1} \\pi^0 \\pi^0$\",\n",
    ")\n",
    "ax.set_xlim(0.25, 2.5)\n",
    "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$ [GeV]\")\n",
    "ax.set_yticks([])\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or generically, so that we can stack all the sub-intensities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "from tensorwaves.data import generate_data\n",
    "\n",
    "logging.basicConfig()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "intensity_components = [\n",
    "    create_intensity_component(model, c, backend=\"numpy\")\n",
    "    for c in model.components\n",
    "    if c.startswith(\"I\")\n",
    "]\n",
    "sub_intensities = [\n",
    "    generate_data(\n",
    "        size=5_000,\n",
    "        reaction_info=model.adapter.reaction_info,\n",
    "        data_transformer=model.adapter,\n",
    "        intensity=component,\n",
    "    )\n",
    "    for component in intensity_components\n",
    "]\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(8, 5))\n",
    "plt.hist(\n",
    "    [model.adapter.transform(i)[\"m_12\"] for i in sub_intensities],\n",
    "    bins=100,\n",
    "    stacked=True,\n",
    "    alpha=0.6,\n",
    ")\n",
    "ax.set_xlim(0.25, 2.5)\n",
    "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$ [GeV]\")\n",
    "ax.set_yticks([])\n",
    "ax.legend(\n",
    "    labels=[f\"${c[2:-1]}$\" for c in model.components if c.startswith(\"I\")]\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyze optimization process"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that {ref}`in Step 3.2 <usage/step3:3.2 Optimize intensity model>`, we initialized {class}`.Minuit2` with a {class}`.TFSummary` callback as well. Its output files provide a nice, interactive representation of the fit process and can be viewed with [TensorBoard](https://www.tensorflow.org/tensorboard/get_started) as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{tabbed} Terminal\n",
    "```bash\n",
    "tensorboard --logdir logs\n",
    "```\n",
    "````\n",
    "\n",
    "````{tabbed} Python\n",
    "```python\n",
    "import tensorboard as tb\n",
    "\n",
    "tb.notebook.list()  # View open TensorBoard instances\n",
    "tb.notebook.start(args_string=\"--logdir logs\")\n",
    "```\n",
    "See more info [here](https://www.tensorflow.org/tensorboard/tensorboard_in_notebooks#tensorboard_in_notebooks)\n",
    "````\n",
    "\n",
    "````{tabbed} Jupyter notebook\n",
    "```ipython\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs\n",
    "```\n",
    "See more info [here](https://www.tensorflow.org/tensorboard/tensorboard_in_notebooks#tensorboard_in_notebooks)\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An alternative would be to use the output of the {class}`.CSVSummary` callback. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import sympy as sp\n",
    "\n",
    "converters = {p: lambda s: complex(s).real for p in initial_parameters}\n",
    "fit_traceback = pd.read_csv(\"fit_traceback_minuit.csv\", converters=converters)\n",
    "fig, (ax1, ax2) = plt.subplots(\n",
    "    2, figsize=(7, 8), gridspec_kw={\"height_ratios\": [1, 1.8]}\n",
    ")\n",
    "fit_traceback.plot(\"function_call\", \"estimator_value\", ax=ax1)\n",
    "fit_traceback.plot(\"function_call\", sorted(initial_parameters), ax=ax2)\n",
    "fig.suptitle(\"SciPy optimizer process\", fontsize=16)\n",
    "ax1.set_title(\"Negative log likelihood\")\n",
    "ax2.set_title(\"Parameter values\")\n",
    "ax1.set_xlabel(\"function call\")\n",
    "ax2.set_xlabel(\"function call\")\n",
    "fig.tight_layout()\n",
    "ax1.legend().remove()\n",
    "legend_texts = ax2.legend().get_texts()\n",
    "for text in legend_texts:\n",
    "    latex = f\"${sp.latex(sp.Symbol(text.get_text()))}$\"\n",
    "    latex = latex.replace(\"\\\\\\\\\", \"\\\\\")\n",
    "    if latex[2] == \"C\":\n",
    "        latex = Rf\"\\left|{latex}\\right|\"\n",
    "    text.set_text(latex)\n",
    "for line in ax2.get_lines():\n",
    "    label = line.get_label()\n",
    "    color = line.get_color()\n",
    "    ax2.axhline(\n",
    "        y=complex(sympy_model.parameters[label]).real,\n",
    "        color=color,\n",
    "        alpha=0.5,\n",
    "        linestyle=\"dotted\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SciPy optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an alternative to the {class}`.Minuit2` optimizer, you can use the {class}`.ScipyMinimizer`. As opposed to Minuit, this optimizer does not compute errors. See {doc}`scipy:tutorial/optimize` and {mod}`scipy.optimize` for more info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.optimizer.scipy import ScipyMinimizer\n",
    "\n",
    "scipy = ScipyMinimizer(\n",
    "    callback=CSVSummary(\n",
    "        \"fit_traceback_scipy.csv\",\n",
    "        function_call_step_size=1,\n",
    "        iteration_step_size=1,\n",
    "    ),\n",
    "    use_analytic_gradient=False,\n",
    ")\n",
    "fit_result = scipy.optimize(estimator, initial_parameters)\n",
    "fit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "converters = {p: lambda s: complex(s).real for p in initial_parameters}\n",
    "fit_traceback = pd.read_csv(\"fit_traceback_scipy.csv\", converters=converters)\n",
    "fig, (ax1, ax2) = plt.subplots(\n",
    "    2, figsize=(7, 8), gridspec_kw={\"height_ratios\": [1, 1.8]}\n",
    ")\n",
    "fit_traceback.plot(\"function_call\", \"estimator_value\", ax=ax1)\n",
    "fit_traceback.plot(\"function_call\", sorted(initial_parameters), ax=ax2)\n",
    "fig.suptitle(\"SciPy optimizer process\", fontsize=16)\n",
    "ax1.set_title(\"Negative log likelihood\")\n",
    "ax2.set_title(\"Parameter values\")\n",
    "ax1.set_xlabel(\"function call\")\n",
    "ax2.set_xlabel(\"function call\")\n",
    "fig.tight_layout()\n",
    "ax1.legend().remove()\n",
    "legend_texts = ax2.legend().get_texts()\n",
    "for text in legend_texts:\n",
    "    latex = f\"${sp.latex(sp.Symbol(text.get_text()))}$\"\n",
    "    latex = latex.replace(\"\\\\\\\\\", \"\\\\\")\n",
    "    if latex[2] == \"C\":\n",
    "        latex = Rf\"\\left|{latex}\\right|\"\n",
    "    text.set_text(latex)\n",
    "for line in ax2.get_lines():\n",
    "    label = line.get_label()\n",
    "    color = line.get_color()\n",
    "    ax2.axhline(\n",
    "        y=complex(sympy_model.parameters[label]).real,\n",
    "        color=color,\n",
    "        alpha=0.5,\n",
    "        linestyle=\"dotted\",\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
