{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Analyze results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "\n",
    "The {doc}`pycompwa </index>` is no longer maintained. Use the [ComPWA](https://compwa-org.rtfd.io) packages [QRules](https://qrules.rtfd.io), [AmpForm](https://ampform.rtfd.io), and [TensorWaves](https://tensorwaves.rtfd.io) instead!\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "We load the output from the previous steps first. Note that you need to define an estimator before loading the fit result to the intensity.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycompwa.ui as pwa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwa.Logging(\"error\")  # at this stage, we are not interested in the back-end\n",
    "particle_list = pwa.read_particles(\"model.xml\")\n",
    "kinematics = pwa.create_helicity_kinematics(\"model.xml\", particle_list)\n",
    "kinematics.create_all_subsystems()\n",
    "data_sample = pwa.read_root_data(input_file=\"generated_data.root\")\n",
    "phsp_sample = pwa.read_root_data(input_file=\"generated_phsp.root\")\n",
    "intensity_builder = pwa.IntensityBuilderXML(\n",
    "    \"model.xml\", particle_list, kinematics, phsp_sample\n",
    ")\n",
    "fit_result = pwa.load(\"fit_result.xml\")\n",
    "intensity = intensity_builder.create_intensity()\n",
    "data_set = kinematics.convert(data_sample)\n",
    "phsp_set = kinematics.convert(phsp_sample)\n",
    "(\n",
    "    estimator,\n",
    "    initial_parameters,\n",
    ") = pwa.create_unbinned_log_likelihood_function_tree_estimator(\n",
    "    intensity, data_set\n",
    ")\n",
    "intensity.updateParametersFrom(fit_result.final_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.1 Convert to pandas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is easiest to analyze your data with {class}`pandas.DataFrame`. We are mostly interested in the distributions of the kinematic variables, so we first have to convert the {class}`.DataSet` instances to a {class}`~pandas.DataFrame`. We can do this with the {mod}`.data.convert` module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.data import convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_data = convert.data_set_to_pandas(data_set)\n",
    "frame_phsp = convert.data_set_to_pandas(phsp_set)\n",
    "frame_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We converted the phase space sample as well, because it serves as a uniform space over which to evaluate the intensity model. We therefore attribute an intensity as a weight to each entry in the phase space {class}`~pandas.DataFrame`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity_set = intensity.evaluate(phsp_set.data)\n",
    "frame_phsp[\"weights\"] = intensity_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{hint}\n",
    "\n",
    "The {mod}`pycompwa.data` contains several tools for converting from and to {class}`pandas.DataFrame`. There is also a {mod}`.data.io` submodule, which allows you to import form and export to file formats of other PWA frameworks. A demonstration of this module can be found under {doc}`/usage/tools/data_module`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 Visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we converted our data to {class}`pandas.DataFrame`, we can use its full potential! In this section, we will use some of that {class}`~pandas.DataFrame` functionality to create histograms of the distributions of the kinematic variables and will have a look at the Dalitz plots both data and fit result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kinematic variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, the columns of the {class}`~pandas.DataFrame` created above contain final state IDs. Of course, you would rather see particle names there, but we need to use unique IDs in case there are identical particle names. The {mod}`.data.naming` module provides a tool that can convert the final state IDs to recognizable names:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.data import naming\n",
    "\n",
    "naming.replace_ids(\"2,3\", \"model.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can apply the usual {mod}`matplotlib.pyplot` tools to plot the distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_1d_comparison(name, bins=100, **kwargs):\n",
    "    \"\"\"Helper function for comparing the 1D distributions of fit and data.\"\"\"\n",
    "    frame_data[name].hist(\n",
    "        bins=bins, density=True, alpha=0.5, label=\"data\", **kwargs\n",
    "    )\n",
    "    frame_phsp[name].hist(\n",
    "        bins=bins,\n",
    "        weights=frame_phsp[\"weights\"],\n",
    "        density=True,\n",
    "        histtype=\"step\",\n",
    "        color=\"red\",\n",
    "        label=\"fit\",\n",
    "        **kwargs,\n",
    "    )\n",
    "    plt.ylabel(\"normalized intensity\")\n",
    "    title = naming.replace_ids(name, kinematics)\n",
    "    plt.xlabel(title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_1d_comparison(\"theta_2_4_vs_3\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Applying a transformation to one of the columns is also easily done for a {class}`pandas.DataFrame`. For instance, if you are interested in the invariant mass (and not the **square value** of the invariant mass), simply use {obj}`numpy.sqrt() <numpy.sqrt>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy import sqrt\n",
    "\n",
    "frame_data[\"M(3,4)\"] = sqrt(frame_data[\"mSq_(3,4)\"])\n",
    "frame_phsp[\"M(3,4)\"] = sqrt(frame_phsp[\"mSq_(3,4)\"])\n",
    "\n",
    "plot_1d_comparison(\"M(3,4)\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dalitz plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dalitz plots are 2-dimensional histograms of the square values of the invariant masses. We are therefore interested in the following variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variable_names = data_set.data.keys()\n",
    "[var for var in variable_names if var.startswith(\"mSq_\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now use {func}`~matplotlib.pyplot.hist2d` to generate the Dalitz plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dalitz_plot(frame, mass_x, mass_y, bins=50, **kwargs):\n",
    "    \"\"\"Helper function to create a Dalitz plot with useful axis titles.\"\"\"\n",
    "    plt.hist2d(frame[mass_x], frame[mass_y], bins=bins, **kwargs)\n",
    "    plt.xlabel(naming.replace_ids(mass_x, kinematics))\n",
    "    plt.ylabel(naming.replace_ids(mass_y, kinematics))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.colors import LogNorm  # logarithmic z-axis\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(9, 4))\n",
    "\n",
    "plt.sca(axs[0])\n",
    "axs[0].set_title(\"data\")\n",
    "dalitz_plot(frame_data, \"mSq_(3,4)\", \"mSq_(2,4)\", norm=LogNorm())\n",
    "\n",
    "plt.sca(axs[1])\n",
    "axs[1].set_title(\"fit\")\n",
    "dalitz_plot(\n",
    "    frame_phsp,\n",
    "    \"mSq_(3,4)\",\n",
    "    \"mSq_(2,4)\",\n",
    "    norm=LogNorm(),\n",
    "    weights=frame_phsp[\"weights\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.3 Calculate fit fractions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original intensity model that we optimized consists of several components. All registered component names can be retrieved via {meth}`.get_all_component_names()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity_builder.get_all_component_names()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, these components are amplitudes and intensities that are added coherently or incoherently.\n",
    "\n",
    "In PWA, you are often interested in the fractions that those amplitudes contribute to the total intensity. These **fit fractions** can be calculated using the {func}`.fit_fractions_with_propagated_errors` function. It requires amplitude or {class}`.Intensity` components that can be extracted from the {class}`.IntensityBuilderXML` instance. A nested list of the component names has to be specified as well. These are the names specified in the component XML attributes. If the inner lists contain more than one component, these components will be added coherently. In this way, you can calculate your own customized fit fractions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "components = intensity_builder.create_intensity_components(\n",
    "    [\n",
    "        [\"coherent_J/psi_-1_to_gamma_-1+pi0_0+pi0_0\"],\n",
    "        [\"J/psi_-1_to_f2(1270)_-1+gamma_-1;f2(1270)_-1_to_pi0_0+pi0_0;\"],\n",
    "        [\"J/psi_-1_to_f2(1270)_-2+gamma_-1;f2(1270)_-2_to_pi0_0+pi0_0;\"],\n",
    "        [\"J/psi_-1_to_f2(1270)_0+gamma_-1;f2(1270)_0_to_pi0_0+pi0_0;\"],\n",
    "        [\n",
    "            \"J/psi_-1_to_f2(1270)_-1+gamma_-1;f2(1270)_-1_to_pi0_0+pi0_0;\",\n",
    "            \"J/psi_-1_to_f2(1270)_-2+gamma_-1;f2(1270)_-2_to_pi0_0+pi0_0;\",\n",
    "            \"J/psi_-1_to_f2(1270)_0+gamma_-1;f2(1270)_0_to_pi0_0+pi0_0;\",\n",
    "        ],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last step is to call the actual fit fraction calculation function {func}`.fit_fractions_with_propagated_errors`. Here, the first required argument is a list of component pairs. Each pair resembles the nominator and denominator of a fit fraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_fractions = pwa.fit_fractions_with_propagated_errors(\n",
    "    [\n",
    "        (components[1], components[0]),\n",
    "        (components[2], components[0]),\n",
    "        (components[3], components[0]),\n",
    "        (components[4], components[0]),\n",
    "    ],\n",
    "    phsp_set,\n",
    "    fit_result,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for fraction in fit_fractions:\n",
    "    print(fraction.name.replace(\";\", \"\\n  \"))\n",
    "    print(fraction.value, \"+/-\", fraction.error)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! You can check some of the other {doc}`examples </usage>` to learn about more detailed features of ComPWA.\n",
    "\n",
    "And, please, feel free to [provide feedback](https://github.com/ComPWA/pycompwa/issues/new) or [contribute](https://compwa.github.io/develop.html) ;)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
