docs/amplitude-analysis.ipynb

{
 "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,pwa,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": [
    "import os\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Amplitude analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While TensorWaves can handle [general mathematical expressions](./usage.ipynb), it was originally created to perform **amplitude analysis** / **Partial Wave Analysis** (PWA), that is, to fit amplitude models to four-momenta data samples.\n",
    "\n",
    "This notebook shows how to do an amplitude analysis with the [ComPWA](https://github.com/ComPWA) packages [QRules](https://qrules.rtfd.io), [AmpForm](https://AmpForm.rtfd.io), and [TensorWaves](https://tensorwaves.rtfd.io). The ComPWA workflow generally consists of three stages:\n",
    "\n",
    "1. [Create an amplitude model](compwa-step-1) with {mod}`qrules` and {mod}`ampform`.\n",
    "2. [Generate hit-and-miss data samples](compwa-step-2) with this amplitude model.\n",
    "3. [Fit model to the data samples](compwa-step-3).\n",
    "\n",
    ":::{note}\n",
    "\n",
    "This notebook shows several tricks that _can_ be helpful when doing an amplitude analysis. **Most steps are optional** though—they only serve to illustrate some tips that can be adopted and worked out further for specific analyses.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "(compwa-step-1)=\n",
    "## Step 1: Formulate model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "Whether [generating data](compwa-step-2) or [fitting a model](compwa-step-3), TensorWaves takes mathematical expressions as input. When that expression is an amplitude model, it is most convenient to formulate it with {mod}`qrules` and {mod}`ampform`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Generate transitions\n",
    "\n",
    "The first step is to generate all allowed transitions with {doc}`QRules <qrules:index>`. In this example, we use the helicity formalism, but you can also use `formalism=\"canonical-helicity\"`. As you can see, we analyze the decay $J/\\psi \\to \\gamma\\pi^0\\pi^0$ here.\n",
    "\n",
    ":::{admonition} Simplified model: $J/\\psi \\to \\gamma f_0$, $f_0 \\rightarrow \\pi^0\\pi^0$\n",
    "---\n",
    "class: dropdown\n",
    "---\n",
    "\n",
    "As [](compwa-step-3) serves to illustrate usage only, we make the amplitude model here a bit simpler by not allowing $\\omega$ resonances (which are narrow and therefore hard to fit). For this reason, we can also limit the {class}`~qrules.settings.InteractionType` to {attr}`~qrules.settings.InteractionType.STRONG`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qrules\n",
    "\n",
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [-1, +1]),\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\"f(0)\"],\n",
    "    allowed_interaction_types=[\"strong\", \"EM\"],\n",
    "    formalism=\"helicity\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import graphviz\n",
    "\n",
    "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "\n",
    "See more advanced examples on {doc}`QRules' usage page <qrules:usage>`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(build-amplitude-model)=\n",
    "### 1.2 Build amplitude model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we use {mod}`ampform` to formulate the {attr}`~qrules.transition.ReactionInfo.transitions` as an amplitude model (here: {class}`~ampform.helicity.HelicityModel`). This can be done with {func}`~ampform.get_builder` and {meth}`~ampform.helicity.HelicityAmplitudeBuilder.formulate`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ampform\n",
    "\n",
    "model_builder = ampform.get_builder(reaction)\n",
    "model_no_dynamics = model_builder.formulate()\n",
    "model_no_dynamics.intensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from ampform.io import aslatex\n",
    "from IPython.display import Math\n",
    "\n",
    "Math(aslatex(model_no_dynamics.amplitudes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Math(aslatex(model_no_dynamics.parameter_defaults))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The heart of the model is a sympy {attr}`~ampform.helicity.HelicityModel.expression` that contains the full description of the intensity model. Note two things:\n",
    "1. The coefficients for the different amplitudes are **complex** valued.\n",
    "2. By default there is no dynamics in the model, so it still has to be specified.\n",
    "\n",
    "We choose to use {func}`~ampform.dynamics.relativistic_breit_wigner_with_ff` as the lineshape for all resonances and use a Blatt-Weisskopf form factor ({func}`~ampform.dynamics.builder.create_non_dynamic_with_ff`) for the production decay. The {meth}`~ampform.helicity.DynamicsSelector.assign` method of the {attr}`~ampform.helicity.HelicityAmplitudeBuilder.dynamics` attribute is a convenience interface for replacing the dynamics for intermediate states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics.builder import (\n",
    "    create_non_dynamic_with_ff,\n",
    "    create_relativistic_breit_wigner_with_ff,\n",
    ")\n",
    "\n",
    "model_builder.dynamics.assign(\"J/psi(1S)\", create_non_dynamic_with_ff)\n",
    "for name in reaction.get_intermediate_particles().names:\n",
    "    model_builder.dynamics.assign(name, create_relativistic_breit_wigner_with_ff)\n",
    "model = model_builder.formulate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's take another look at the parameters of the model to see which new parameters are there:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "sorted_parameter_defaults = {\n",
    "    symbol: model.parameter_defaults[symbol]\n",
    "    for symbol in sorted(model.parameter_defaults, key=str)\n",
    "}\n",
    "src = aslatex(sorted_parameter_defaults)\n",
    "Math(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally, we can backup the {class}`~ampform.helicity.HelicityModel` to disk via {mod}`pickle`. The {class}`~qrules.transition.ReactionInfo` object (which takes longest to generate) can also be pickled, or it can also be serialized to JSON:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "qrules.io.write(reaction, \"transitions.json\")\n",
    "with open(\"helicity_model.pickle\", \"wb\") as stream:\n",
    "    pickle.dump(model, stream)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next steps, we use this {class}`~ampform.helicity.HelicityModel` as a template for a computational function to [generate data](compwa-step-2) and to [perform a fit](compwa-step-3).\n",
    "\n",
    ":::{tip}\n",
    "\n",
    "See more advanced examples on {doc}`AmpForm's usage page <ampform:usage>`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "(compwa-step-2)=\n",
    "## Step 2: Generate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "::::{margin}\n",
    ":::{tip}\n",
    "{doc}`TR-018<compwa:report/018>` explains some of the mechanisms behind the phase space generator as well as how to do {ref}`importance sampling<compwa:report/018:Intensity distribution>`.\n",
    ":::\n",
    "::::\n",
    "\n",
    "In this section, we use the {class}`~ampform.helicity.HelicityModel` that we created with {mod}`ampform` in [the previous step](compwa-step-1) to generate a data sample via hit & miss Monte Carlo. We do this with the {mod}`.data` module.\n",
    "\n",
    "**Optionally**, we can {func}`~pickle.load` the {class}`~ampform.helicity.HelicityModel` that was created in the previous step. This does not have to be done if the model has been generated in the same script or notebook (like in this notebook), but can be useful if the model was generated elsewhere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "from ampform.helicity import HelicityModel\n",
    "\n",
    "with open(\"helicity_model.pickle\", \"rb\") as model_file:\n",
    "    imported_model: HelicityModel = pickle.load(model_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "initial_state, *_ = imported_model.reaction_info.initial_state.values()\n",
    "print(\"Initial state:\")\n",
    "print(\" \", initial_state.name)\n",
    "print(\"Final state:\")\n",
    "for i, p in imported_model.reaction_info.final_state.items():\n",
    "    print(f\"  {i}: {p.name}\")\n",
    "del initial_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-step-2.1)=\n",
    "### 2.1 Generate phase space sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {class}`~qrules.transition.ReactionInfo` class defines the constraints of the phase space. As such, we have enough information to generate a **phase-space sample** for this particle reaction. We do this with a {class}`.TFPhaseSpaceGenerator` class, which is a {class}`.DataGenerator` for a {obj}`.DataSample` of **four-momenta** arrays (using {obj}`tensorflow <tf.Tensor>` and the [`phasespace`](https://phasespace.readthedocs.io) package as a back-end). We also need to construct a {class}`.RealNumberGenerator` that can generate random numbers. {class}`.TFUniformRealNumberGenerator` is the natural choice here.\n",
    "\n",
    "As opposed to the main [](compwa-step-2) of the main usage example page, we will generate a **deterministic** data sample. This can be done by feeding a {class}`.RealNumberGenerator` with a specific {attr}`~.RealNumberGenerator.seed` and giving that generator to the {meth}`.TFPhaseSpaceGenerator.generate` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator\n",
    "\n",
    "rng = TFUniformRealNumberGenerator(seed=0)\n",
    "phsp_generator = TFPhaseSpaceGenerator(\n",
    "    initial_state_mass=reaction.initial_state[-1].mass,\n",
    "    final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n",
    ")\n",
    "phsp_momenta = phsp_generator.generate(100_000, rng)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "pd.DataFrame({\n",
    "    (k, label): np.transpose(v)[i]\n",
    "    for k, v in phsp_momenta.items()\n",
    "    for i, label in enumerate([\"E\", \"px\", \"py\", \"pz\"])\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "The resulting phase space sample is a {obj}`dict` of final state IDs to an {obj}`~numpy.array` of four-momenta. In the last step, we converted this sample in such a way that it is rendered as an understandable {class}`pandas.DataFrame`.\n",
    "\n",
    ":::{admonition} Four-momentum arrays\n",
    "{doc}`Kinematic expressions <ampform:usage/kinematics>` from AmpForm that involve four-momenta should be formatted as $\\left(E, p_x, p_y, p_z\\right)$. In addition, the shape of input arrays should be `(n, 4)` with `n` the number of events.\n",
    ":::\n",
    "\n",
    ":::{warning}\n",
    "When using the helicity formalism, the sum of the four-momenta should be in the rest frame, that is $\\sum_i p_i = \\left(m_A, 0, 0, 0\\right)$ with $m_A$ the mass of the decaying particle&nbsp;$A$. This is because the helicity formalisms boosts through the decay chain starting from particle&nbsp;$A$. **Take care to boost your experimental data into the rest frame**, optionally following the {doc}`kinematics classes provided by AmpForm <ampform:usage/kinematics>`.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "p = np.array(list(phsp_momenta.values()))\n",
    "p.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "p.sum(axis=0).round(decimals=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-step-2.2)=\n",
    "### 2.2 Generate intensity-based sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "'Data samples' are more complicated than phase space samples in that they represent the intensity profile resulting from a reaction. You therefore need a {class}`.Function` object that expresses an intensity distribution as well as a phase space over which to generate that distribution. We call such a data sample an **intensity-based sample**.\n",
    "\n",
    "An intensity-based sample is generated over a phase space sample using the {class}`.IntensityDistributionGenerator`. Its usage is similar to {class}`.TFPhaseSpaceGenerator`, but now you have to provide a {obj}`.Function` as well as a {obj}`.DataTransformer` that is used to transform the four-momentum phase space sample to a data sample that can be understood by the {obj}`.Function`.\n",
    "\n",
    "Now, recall that in [](compwa-step-1), we used the helicity formalism to mathematically express the reaction in terms of an amplitude model. TensorWaves needs to convert this {obj}`~ampform.helicity.HelicityModel` to a {obj}`.Function` object that can perform fast computations. This can be done with {func}`.create_parametrized_function`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.function.sympy import create_parametrized_function\n",
    "\n",
    "unfolded_expression = model.expression.doit()\n",
    "intensity_func = create_parametrized_function(\n",
    "    expression=unfolded_expression,\n",
    "    parameters=model.parameter_defaults,\n",
    "    backend=\"numpy\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "\n",
    "If {func}`.create_parametrized_function` takes a long time, have a look at {doc}`usage/faster-lambdify`.\n",
    "\n",
    ":::\n",
    "\n",
    ":::{seealso}\n",
    "\n",
    "{ref}`usage/basics:Hit & miss`\n",
    "\n",
    ":::\n",
    "\n",
    "A problem is that {class}`.ParametrizedBackendFunction` takes a {obj}`.DataSample` with kinematic variables for the helicity formalism as input, not a set of four-momenta. We therefore need to construct a {class}`.DataTransformer` to transform these four-momenta to function variables. In this case, we work with the helicity formalism, so we construct a {class}`.SympyDataTransformer`. This **numerical** data transformer is created from symbolic expressions, in this case, from expressions that relate four-momenta to helicity angles and invariant masses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Math(aslatex(model.kinematic_variables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.data import SympyDataTransformer\n",
    "\n",
    "helicity_transformer = SympyDataTransformer.from_sympy(\n",
    "    model.kinematic_variables, backend=\"jax\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{ref}`usage:Generate and transform data`\n",
    "\n",
    ":::\n",
    "\n",
    "That's it, now we have enough info to create an intensity-based data sample. Notice how the structure of the output data is the same as the [phase-space sample we generated previously](compwa-step-2.1):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "from tensorwaves.data import (\n",
    "    IntensityDistributionGenerator,\n",
    "    TFWeightedPhaseSpaceGenerator,\n",
    ")\n",
    "\n",
    "weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n",
    "    initial_state_mass=reaction.initial_state[-1].mass,\n",
    "    final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n",
    ")\n",
    "data_generator = IntensityDistributionGenerator(\n",
    "    domain_generator=weighted_phsp_generator,\n",
    "    function=intensity_func,\n",
    "    domain_transformer=helicity_transformer,\n",
    ")\n",
    "data_momenta = data_generator.generate(10_000, rng)\n",
    "pd.DataFrame({\n",
    "    (k, label): np.transpose(v)[i]\n",
    "    for k, v in data_momenta.items()\n",
    "    for i, label in enumerate([\"E\", \"px\", \"py\", \"pz\"])\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we use a {class}`.RealNumberGenerator` with a specific {attr}`~.RealNumberGenerator.seed` to ensure we get a **deterministic** data sample.\n",
    "\n",
    ":::{note}\n",
    "\n",
    "Instead of constructing a {class}`.TFWeightedPhaseSpaceGenerator`, it's also possible to just reuse the **unweighted** {class}`.TFPhaseSpaceGenerator` that we constructed previously. This is, however, a bit slower, because the {class}`.TFPhaseSpaceGenerator` also uses a hit-and-miss strategy.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-step-2.3)=\n",
    "### 2.3 Visualize kinematic variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have a phase space sample and an intensity-based sample. Their data structure isn't the most informative though: it's just a collection of four-momentum tuples. But we can again use the {class}`.SympyDataTransformer` to convert these four-momenta to (in the case of the helicity formalism) invariant masses and helicity angles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phsp = helicity_transformer(phsp_momenta)\n",
    "data = helicity_transformer(data_momenta)\n",
    "list(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "::::{note}\n",
    "Check the remark about four-momentum format and rest frame of the decaying particle [here](compwa-step-2.1).\n",
    "::::\n",
    "\n",
    "The {obj}`.DataSample` is a mapping of kinematic variables names to a 1-dimensional array of values. The numbers you see here are final state IDs as defined in the {class}`~ampform.helicity.HelicityModel` member of the {class}`~ampform.helicity.HelicityModel`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "for state_id, particle in reaction.final_state.items():\n",
    "    print(f\"ID {state_id}:\", particle.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Available kinematic variables\n",
    "---\n",
    "class: dropdown\n",
    "---\n",
    "By default, {mod}`tensorwaves` only generates invariant masses of the {class}`Topologies <qrules.topology.Topology>` that are of relevance to the decay problem. In this case, we only have resonances $f_0 \\to \\pi^0\\pi^0$. If you are interested in more invariant mass combinations, you can do so with the method {meth}`~ampform.kinematics.HelicityAdapter.register_topology`.\n",
    "````\n",
    "\n",
    "The {obj}`.DataSample` can easily be converted to a {class}`pandas.DataFrame`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "data_frame = pd.DataFrame(data)\n",
    "phsp_frame = pd.DataFrame(phsp)\n",
    "data_frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This also means that we can use all kinds of fancy plotting functionality of for instance {mod}`matplotlib.pyplot` to see what's going on. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "resonances = sorted(\n",
    "    reaction.get_intermediate_particles(),\n",
    "    key=lambda p: p.mass,\n",
    ")\n",
    "evenly_spaced_interval = np.linspace(0, 1, len(resonances))\n",
    "colors = [plt.cm.rainbow(x) for x in evenly_spaced_interval]\n",
    "fig, ax = plt.subplots(figsize=(9, 4))\n",
    "ax.hist(\n",
    "    np.real(data_frame[\"m_12\"]),\n",
    "    bins=100,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    ")\n",
    "ax.set_xlabel(\"$m$ [GeV]\")\n",
    "for p, color in zip(resonances, colors):\n",
    "    ax.axvline(x=p.mass, linestyle=\"dotted\", label=p.name, color=color)\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "[](#fit-fractions)\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Export data sets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To export the generated data samples, simply {func}`pickle.dump` them as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"data.pickle\", \"wb\") as stream:\n",
    "    pickle.dump(data, stream)\n",
    "with open(\"phsp.pickle\", \"wb\") as stream:\n",
    "    pickle.dump(phsp, stream)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the [next step](compwa-step-3), we illustrate how to {meth}`~.Minuit2.optimize` the intensity model to these data samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "(compwa-step-3)=\n",
    "## Step 3: Perform fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As explained in the [previous step](compwa-step-2), a {class}`.ParametrizedFunction` can compute a list of intensities (real numbers) for an input {obj}`.DataSample`. At this stage, we want to optimize the parameters of this {class}`.ParametrizedFunction`, so that it matches the distribution of our data sample. This is what we call 'fitting'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "reaction = qrules.io.load(\"transitions.json\")\n",
    "with open(\"helicity_model.pickle\", \"rb\") as stream:\n",
    "    model: HelicityModel = pickle.load(stream)\n",
    "with open(\"data.pickle\", \"rb\") as stream:\n",
    "    data = pickle.load(stream)\n",
    "with open(\"phsp.pickle\", \"rb\") as stream:\n",
    "    phsp = pickle.load(stream)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-step-3.1)=\n",
    "### 3.1 Prepare parametrized function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In principle, we can use the same {class}`.ParametrizedFunction` as the one that we created in [](compwa-step-2.2). However, when fitting such a function to a data distribution, an {class}`.Optimizer` will evaluate this function numerous times, so it is smart to apply some optimizations to the underlying expression tree before-hand.\n",
    "\n",
    ":::{tip}\n",
    "\n",
    "The sections below illustrate some tricks for how to simplify the expression tree underneath a {class}`.ParametrizedFunction`. **Most of this can also be achieved with {func}`.create_cached_function`**, which is illustrated in {doc}`/usage/caching` and under [](compwa-create_cached_function). But note that it is still smart to [cast complex-valued data](#cast-complex-valued-data).\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Determine free parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It often happens that not all parameters in a {class}`.ParametrizedFunction` have to be optimized. These parameters are called **fixed parameters**, while the ones that will be optimized are called **free parameters**. In the fit example here, we will optimize the following free parameters, starting with a certain initial value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "initial_parameters = {\n",
    "    R\"C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}\": (\n",
    "        1.0 + 0.0j\n",
    "    ),\n",
    "    \"m_{f_{0}(500)}\": 0.4,\n",
    "    \"m_{f_{0}(980)}\": 0.88,\n",
    "    \"m_{f_{0}(1370)}\": 1.22,\n",
    "    \"m_{f_{0}(1500)}\": 1.45,\n",
    "    \"m_{f_{0}(1710)}\": 1.83,\n",
    "    R\"\\Gamma_{f_{0}(500)}\": 0.3,\n",
    "    R\"\\Gamma_{f_{0}(980)}\": 0.1,\n",
    "    R\"\\Gamma_{f_{0}(1710)}\": 0.3,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "parameter_names = {symbol.name for symbol in model.parameter_defaults}\n",
    "not_in_model = {\n",
    "    par_name for par_name in initial_parameters if par_name not in parameter_names\n",
    "}\n",
    "if len(not_in_model) != 0:\n",
    "    msg = f\"Parameters {', '.join(sorted(not_in_model))} do not exist in model\"\n",
    "    raise ValueError(msg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make the fit more interesting, we give these initial values a small offset compared to the suggested {attr}`~ampform.helicity.HelicityModel.parameter_defaults`―after all, we are fitting to a data sample that was generated with this very same {class}`.ParametrizedFunction`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{admonition} Complex-valued parameters\n",
    "\n",
    "If initial parameter values are {obj}`complex`, the parameter is split into a real and an imaginary part during the fit. See also [](#covariance-matrix).\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Collapsing constant sub-trees"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having decided which parameters are to be optimized, we can apply some additional optimizations to the amplitude model expression, so that the computations during the fit are faster. The first of these optimizations is to substitute the parameters that remain fixed with their suggested parameter values. Note how this decreases the number of operations (nodes) in the expression tree:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "\n",
    "sp.count_ops(unfolded_expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "free_parameters = {p for p in model.parameter_defaults if p.name in initial_parameters}\n",
    "fixed_parameters = {\n",
    "    p: v for p, v in model.parameter_defaults.items() if p not in free_parameters\n",
    "}\n",
    "optimized_expression = unfolded_expression.subs(fixed_parameters)\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "old = sp.count_ops(unfolded_expression)\n",
    "new = sp.count_ops(optimized_expression)\n",
    "assert old > new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that there are a few irrational numbers (square roots) in the expression as well, for example in this amplitude:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_expression.args[0].args[0].args[0].args[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we substitute these values, the expression tree becomes even smaller:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for node in sp.preorder_traversal(optimized_expression):\n",
    "    if node.free_symbols:\n",
    "        continue\n",
    "    if isinstance(node, sp.Pow):\n",
    "        optimized_expression = optimized_expression.xreplace({node: node.n()})\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "old = sp.count_ops(unfolded_expression)\n",
    "new = sp.count_ops(optimized_expression)\n",
    "assert old > new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Substitute scalar masses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the final state particles $\\gamma, \\pi^0, \\pi^0$ are **stable**, we can substitute their invariant masses $m_0, m_1, m_2$ with scalar values. SymPy will then further simply the expression, so that the computation is less complex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mass_substitutions = {\n",
    "    sp.Symbol(f\"m_{i}\", nonnegative=True): particle.mass\n",
    "    for i, particle in reaction.final_state.items()\n",
    "}\n",
    "optimized_expression = optimized_expression.subs(mass_substitutions)\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "old = sp.count_ops(unfolded_expression)\n",
    "new = sp.count_ops(optimized_expression)\n",
    "assert old > new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Cast complex-valued data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that some of the computed kinematic variable arrays are complex-valued. This can be useful in come cases, but in this decay channel, all kinematic variable values lie on the real axis. The fit can therefore be further optimized by casting the data arrays to real values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.interface import DataSample\n",
    "\n",
    "\n",
    "def safe_downcast_to_real(data: DataSample) -> DataSample:\n",
    "    # using isrealobj instead of real_if_close to keep same array backend\n",
    "    return {k: v.real if np.isrealobj(v) else v for k, v in data.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_real = safe_downcast_to_real(data)\n",
    "phsp_real = safe_downcast_to_real(phsp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create computational function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we again use {func}`.create_parametrized_function` to create a {class}`.ParametrizedFunction` for this **optimized** expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_function = create_parametrized_function(\n",
    "    optimized_expression, model.parameter_defaults, backend=\"jax\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the intensities computed by the optimized function are indeed the same as the original intensity function that was created in [](compwa-step-2.2) and that it is much faster!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# JIT-compile functions and test equality\n",
    "np.testing.assert_array_almost_equal(\n",
    "    optimized_function(data_real),\n",
    "    intensity_func(data_real),\n",
    "    decimal=13,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n1 intensity_func(data)\n",
    "\n",
    "%timeit -n1 optimized_function(data_real)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reason is that several sub-trees in the original expression tree have been collapsed, which results in fewer computations in the backend. Here's one of the sub-amplitudes, taken from the total expression:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Original expression tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = sp.dotprint(unfolded_expression.args[0].args[0].args[0].args[0])\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Optimized expression tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = sp.dotprint(optimized_expression.args[0].args[0].args[0].args[0])\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-create_cached_function)=\n",
    "#### Simplified procedure: {func}`.create_cached_function`\n",
    "\n",
    "As noted under [](compwa-step-3.1), most of what is described in this section can be achieved with the function {func}`.create_cached_function`. The idea is described on {doc}`/usage/caching`, but in this section, we show how this translates to amplitude analysis.\n",
    "\n",
    "First, note that {func}`.create_cached_function` works with mappings and iterable of {class}`sympy.Symbol <sympy.core.symbol.Symbol>` and not with the {obj}`str` mappings that we defined in [](#determine-free-parameters). We can convert the parameter names back to {class}`~sympy.core.symbol.Symbol`s as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "free_parameter_symbols = [\n",
    "    symbol\n",
    "    for symbol in model.parameter_defaults\n",
    "    if symbol.name in set(initial_parameters)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert len(free_parameter_symbols) == len(initial_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This gives us all the information we need to create a cached function, convert the data samples to cached data samples and construct an {class}`.Estimator` with these transformed items (compare [](compwa-step-3.2)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.estimator import UnbinnedNLL, create_cached_function\n",
    "\n",
    "cached_function, cache_transformer = create_cached_function(\n",
    "    unfolded_expression,\n",
    "    parameters=model.parameter_defaults,\n",
    "    free_parameters=free_parameter_symbols,\n",
    "    backend=\"jax\",\n",
    ")\n",
    "cached_data = cache_transformer(data_real)\n",
    "cached_phsp = cache_transformer(phsp_real)\n",
    "estimator_with_caching = UnbinnedNLL(\n",
    "    cached_function,\n",
    "    data=cached_data,\n",
    "    phsp=cached_phsp,\n",
    "    backend=\"jax\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that, just like in [](#create-computational-function), the computed intensities of both the original intensity function and the cached function are indeed the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_array_almost_equal(\n",
    "    cached_function(cached_data),\n",
    "    intensity_func(data_real),\n",
    "    decimal=13,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n1 intensity_func(data_real)\n",
    "\n",
    "%timeit -n1 cached_function(cached_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-step-3.2)=\n",
    "### 3.2 Define estimator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform a fit, you need to define an {class}`.Estimator`. This is a measure for the discrepancy between the {class}`.ParametrizedFunction` and the data distribution to which you fit it. In PWA, we usually use an **unbinned negative log likelihood estimator** ({class}`.UnbinnedNLL`).\n",
    "\n",
    "Generally, the {class}`.ParametrizedFunction` is not normalized with regards to the data sample, while a log likelihood estimator requires a normalized function. This is where the [phase space data](compwa-step-2.1) comes into play again: the {class}`.ParametrizedFunction` is evaluated over the phase space data, so that its output can be used as a normalization factor.\n",
    "\n",
    "```{margin}\n",
    "If you want to correct for the efficiency of the detector, you should use a *detector-reconstructed* phase space sample.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.estimator import UnbinnedNLL\n",
    "\n",
    "estimator = UnbinnedNLL(\n",
    "    optimized_function,\n",
    "    data=data_real,\n",
    "    phsp=phsp_real,\n",
    "    backend=\"jax\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the {class}`.UnbinnedNLL` can be expressed with different backends, because it uses statistical operations like `log` and `mean`. Here, we use {func}`jax <jax.jit>`, which turns out to be the fastest backend for this model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(compwa-step-3.3)=\n",
    "### 3.3 Optimize fit parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting the fit itself is quite simple: just create an {mod}`.optimizer` instance of your choice and call its {meth}`~.Optimizer.optimize` method to start the fitting process. The {meth}`~.Optimizer.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 have a look at our [first guess for the parameter values](#determine-free-parameters). Recall that a {class}`.ParametrizedFunction` object computes the intensity for a certain {obj}`.DataSample`. This can be seen nicely when we use these intensities as weights on the phase space sample and plot it together with the original data sample. Here, we look at the invariant mass distribution projection of the final states `1` and `2`, which, [as we saw before](compwa-step-2.3), is the final state particle pair $\\pi^0\\pi^0$.\n",
    "\n",
    "Don't forget to use {meth}`~.ParametrizedFunction.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",
    "\n",
    "resonances = sorted(\n",
    "    reaction.get_intermediate_particles(),\n",
    "    key=lambda p: p.mass,\n",
    ")\n",
    "\n",
    "evenly_spaced_interval = np.linspace(0, 1, len(resonances))\n",
    "colors = [plt.cm.rainbow(x) for x in evenly_spaced_interval]\n",
    "\n",
    "\n",
    "def indicate_masses(ax):\n",
    "    ax.set_xlabel(\"$m$ [GeV]\")\n",
    "    for color, resonance in zip(colors, resonances):\n",
    "        ax.axvline(\n",
    "            x=resonance.mass,\n",
    "            linestyle=\"dotted\",\n",
    "            label=resonance.name,\n",
    "            color=color,\n",
    "        )\n",
    "\n",
    "\n",
    "def compare_model(\n",
    "    variable_name,\n",
    "    data,\n",
    "    phsp,\n",
    "    function,\n",
    "    bins=100,\n",
    "):\n",
    "    intensities = function(phsp)\n",
    "    _, ax = plt.subplots(figsize=(9, 4))\n",
    "    data_projection = np.real(data[variable_name])\n",
    "    ax = plt.gca()\n",
    "    ax.hist(\n",
    "        data_projection,\n",
    "        bins=bins,\n",
    "        alpha=0.5,\n",
    "        label=\"data\",\n",
    "        density=True,\n",
    "    )\n",
    "    phsp_projection = np.real(phsp[variable_name])\n",
    "    ax.hist(\n",
    "        phsp_projection,\n",
    "        weights=np.array(intensities),\n",
    "        bins=bins,\n",
    "        histtype=\"step\",\n",
    "        color=\"red\",\n",
    "        label=\"fit model\",\n",
    "        density=True,\n",
    "    )\n",
    "    indicate_masses(ax)\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_parameters = optimized_function.parameters\n",
    "optimized_function.update_parameters(initial_parameters)\n",
    "compare_model(\"m_12\", data_real, phsp_real, optimized_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we are ready to create an {class}`.Optimizer` that can {meth}`~.Minuit2.optimize` the parameters in the {class}`.ParametrizedFunction` (which is embedded in the {class}`.Estimator`). Here, we choose the {class}`.Minuit2` optimizer, which is the most common optimizer in high-energy physics (see also {ref}`usage/basics:Perform fit with different optimizers`). Since we constructed the {class}`.UnbinnedNLL` with {obj}`jax <jax.jit>`, can an optimize the model with analytic gradient over the {class}`.ParametrizedBackendFunction`:\n",
    "\n",
    ":::{note}\n",
    "\n",
    "The computation time depends on the complexity of the model, the 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 minute.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "from tensorwaves.optimizer import Minuit2\n",
    "from tensorwaves.optimizer.callbacks import CSVSummary\n",
    "\n",
    "minuit2 = Minuit2(\n",
    "    callback=CSVSummary(\"fit_traceback.csv\"),\n",
    "    use_analytic_gradient=False,\n",
    ")\n",
    "fit_result = minuit2.optimize(estimator, initial_parameters)\n",
    "fit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert fit_result.minimum_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "See {ref}`usage/basics:Minuit2` for how to tweak the internal {class}`iminuit.Minuit` optimizer.\n",
    "\n",
    ":::\n",
    "\n",
    "As can be seen, the values of the optimized parameters in the {class}`.FitResult` are again comparable to the original parameter values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "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:  {original_parameters[p]:.3}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.4 Export and import"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In [](compwa-step-3.3), we initialized {obj}`.Minuit2` with a {class}`.Loadable` callback. 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": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "latest_parameters = CSVSummary.load_latest_parameters(\"fit_traceback.csv\")\n",
    "latest_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso} \n",
    "\n",
    "{ref}`usage/basics:Callbacks` and [this example](./usage.ipynb#optimize-parameters) of a (custom) plotting callback.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.5 Analyze fit result"
   ]
  },
  {
   "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}`.ParametrizedFunction` and plot it again over the phase space sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_function.update_parameters(fit_result.parameter_values)\n",
    "compare_model(\"m_12\", data_real, phsp_real, optimized_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Covariance matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of the {mod}`.optimizer`s offer more specific information about the fit result. This information can be accessed with {attr}`.FitResult.specifics`. A common example would be to get the {attr}`~iminuit.Minuit.covariance` matrix and {attr}`~iminuit.util.Matrix.correlation` matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covariance_matrix = fit_result.specifics.covariance\n",
    "covariance_matrix.correlation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "Optimizers can only work with real numbers. For this reason, {obj}`complex`-valued parameters are split into a real and an imaginary part. This becomes apparent if we look at the covariance matrix.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AIC and BIC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [AIC](https://en.wikipedia.org/wiki/Akaike_information_criterion#Definition) and [BIC](https://en.wikipedia.org/wiki/Bayesian_information_criterion#Definition) give us another quantification of the quality of the fit. Here's how to compute them from this {obj}`.FitResult`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_real_par = fit_result.count_number_of_parameters(complex_twice=True)\n",
    "n_events = len(next(iter(data.values())))\n",
    "log_likelihood = -fit_result.estimator_value\n",
    "\n",
    "aic = 2 * n_real_par - 2 * log_likelihood\n",
    "bic = n_real_par * np.log(n_events) - 2 * log_likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "print(\"AIC:\", aic)\n",
    "print(\"BIC:\", bic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Fit fractions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we have seen [when formulating the amplitude model](#build-amplitude-model), the helicity model is built up of an incoherent sum of real-valued intensities (one for each spin combination), which are each built up of a coherent sum of complex-valued amplitudes (one for each resonance). If we want to compute what each of these amplitudes contribute to the main intensity, we should use the optimized intensity and set coefficients or helicity couplings of amplitudes that were are not interested in to zero."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example function that can do this. Using regular expressions, we set all coefficients in the intensity function to zero if they do not contain a certain resonance $\\LaTeX$ name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "from tensorwaves.interface import ParametrizedFunction\n",
    "\n",
    "\n",
    "def compute_sub_intensity(\n",
    "    func: ParametrizedFunction,\n",
    "    input_data: DataSample,\n",
    "    resonances: list[str],\n",
    "):\n",
    "    original_parameters = dict(func.parameters)\n",
    "    negative_lookahead = f\"(?!{'|'.join(map(re.escape, resonances))})\"\n",
    "    # https://regex101.com/r/WrgGyD/1\n",
    "    pattern = rf\"^(\\\\mathcal{{H}}|C_)({negative_lookahead}.)*$\"\n",
    "    set_parameters_to_zero(func, pattern)\n",
    "    array = func(input_data)\n",
    "    func.update_parameters(original_parameters)\n",
    "    return array\n",
    "\n",
    "\n",
    "def set_parameters_to_zero(func: ParametrizedFunction, name_pattern: str) -> None:\n",
    "    new_parameters = dict(func.parameters)\n",
    "    for par_name in func.parameters:\n",
    "        if re.match(name_pattern, par_name) is not None:\n",
    "            new_parameters[par_name] = 0\n",
    "    func.update_parameters(new_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "all_resonances_latex = [p.latex for p in resonances]\n",
    "np.testing.assert_array_equal(\n",
    "    compute_sub_intensity(intensity_func, phsp, all_resonances_latex),\n",
    "    intensity_func(phsp),\n",
    ")\n",
    "del all_resonances_latex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These functions can be used to compute and visualize the sub-intensity distributions over the phase space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_intensities = intensity_func(phsp)\n",
    "sub_intensities = {\n",
    "    p: compute_sub_intensity(intensity_func, phsp, resonances=[p.latex])\n",
    "    for p in resonances\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "ax.set_xlim(0.25, 2.5)\n",
    "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$ [GeV]\")\n",
    "ax.set_yticks([])\n",
    "\n",
    "bins = 150\n",
    "phsp_projection = phsp[\"m_12\"].real\n",
    "ax.hist(\n",
    "    phsp_projection,\n",
    "    weights=total_intensities,\n",
    "    bins=bins,\n",
    "    color=\"red\",\n",
    "    histtype=\"step\",\n",
    "    label=\"full intensity\",\n",
    ")\n",
    "ax.hist(\n",
    "    len(sub_intensities) * [phsp_projection],\n",
    "    weights=list(sub_intensities.values()),\n",
    "    bins=bins,\n",
    "    alpha=0.6,\n",
    "    label=[Rf\"${p.latex}$\" for p in sub_intensities],\n",
    "    stacked=True,\n",
    ")\n",
    "\n",
    "fig.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use these functions to compute the decay rates for each resonance. Notice how the sum of the decay rates does not add up to a 100%. This is because of the strong constructive interference between the resonances in this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_intensity = intensity_func(phsp).sum()\n",
    "fit_fractions = {\n",
    "    resonance.name: f\"{sub_intensity.sum() / total_intensity:.1%}\"\n",
    "    for resonance, sub_intensity in sub_intensities.items()\n",
    "}\n",
    "fit_fractions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced examples\n",
    "\n",
    "```{toctree}\n",
    "---\n",
    "maxdepth: 2\n",
    "---\n",
    "amplitude-analysis/analytic-continuation\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "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",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}