{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data module"
   ]
  },
  {
   "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": [
    "The {mod}`pycompwa.data` module provides several tools for importing, exporting, and visualizing data. By data, we mean *event-wise collections of four-momentum tuples*, possibly organized by particle name. We choose to work with {mod}`pandas` as a back-end, because it allows fast manipulation and visualization of data sets and can import and export to several standard data formats.\n",
    "\n",
    "This notebook shows how to conveniently manipulate momentum tuple collections in such a way that they can be imported into ComPWA. It is also shown how to import data from other frameworks and how to do conversions to kinematic variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PWA data frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Input data for PWA frameworks mainly consists of event-wise four-momentum tuples, grouped by particle. The core of the {mod}`.data` module is therefore handled by a specially formatted {class}`pandas.DataFrame`. Such as specific format not only allows us to import and export to different file formats, but also to convert the data to ComPWA objects, such as an {class}`.EventCollection`.\n",
    "\n",
    "The format is guaranteed through a decorator called {func}`~pandas.api.extensions.register_dataframe_accessor`. Such an accessor extends a {class}`~pandas.DataFrame` with several properties, granted that the {class}`~pandas.DataFrame` properly validates according to that accessor. In the {mod}`.data` module, this accessor is the {class}`.PwaAccessor`. We call a {class}`~pandas.DataFrame` that is formatted according to this accessor a {class}`PWA DataFrame <.PwaAccessor>`.\n",
    "\n",
    "To be sure, all this sounds a bit abstract. To illustrate the usage of this accessor, we therefore first have to create a skeleton {class}`PWA DataFrame <.PwaAccessor>`. This can be done through the {mod}`.create` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.data import create"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame = create.pwa_frame(\n",
    "    particle_names=[\"gamma\", \"pi0\", \"pi0\"], number_of_rows=3\n",
    ")\n",
    "frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that his {class}`~pandas.DataFrame` has hierarchical column name (see [multi-indexing](https://pandas.pydata.org/pandas-docs/stable/user_guide/advanced.html)): the first column layer is the particle name, the second contains the four-momentum labels. In addition, duplicate particle names have been made unique by adding an index. Values are undefined by default, but you can set them later on. Here, we do this manually, but you can use this procedure for importing large data sets in your own Python scripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame[\"gamma\", \"p_x\"] = [-0.520903, -0.285015, 0.632325]\n",
    "frame[\"gamma\", \"p_y\"] = [0.885259, 0.520381, -0.779928]\n",
    "frame[\"gamma\", \"p_z\"] = [0.655934, -0.996574, -0.892786]\n",
    "frame[\"gamma\", \"E\"] = [1.21872, 1.15982, 1.34357]\n",
    "\n",
    "frame[\"pi0-1\", \"p_x\"] = [0.653672, 0.452265, 0.113717]\n",
    "frame[\"pi0-1\", \"p_y\"] = [-0.813022, -0.76188, 0.605441]\n",
    "frame[\"pi0-1\", \"p_z\"] = [-1.01763, 0.00723327, 0.718613]\n",
    "frame[\"pi0-1\", \"E\"] = [1.46359, 0.896256, 0.956093]\n",
    "\n",
    "frame[\"pi0-2\", \"p_x\"] = [-0.132769, -0.16725, -0.746043]\n",
    "frame[\"pi0-2\", \"p_y\"] = [-0.0722372, 0.241499, 0.174487]\n",
    "frame[\"pi0-2\", \"p_z\"] = [0.361697, 0.989341, 0.174172]\n",
    "frame[\"pi0-2\", \"E\"] = [0.414596, 1.04082, 0.797233]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now already have a glance at some of the properties that the {class}`.PwaAccessor` offers. You can access these properties through the `pwa` namespace and perform some standard {mod}`pandas` computation on them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Particles:\", frame.pwa.particles)\n",
    "print(\"Momentum labels:\", frame.pwa.momentum_labels)\n",
    "print(\"Weights:\", frame.pwa.weights)\n",
    "print(\"Average pi0 mass:\\n\", frame[[\"pi0-1\", \"pi0-2\"]].pwa.mass.stack().mean())\n",
    "print(\n",
    "    \"Average gamma 3-momentum:\\n\",\n",
    "    frame[\"gamma\"].pwa.rho.mean(),\n",
    "    \"+/-\",\n",
    "    frame[\"gamma\"].pwa.rho.std(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll see more of these properties after we import some real data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and export data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The module {mod}`.data.io` allows one to import from and to data formats of other PWA frameworks. Here's an example, importing a `pawianHists.root` file. Such a file not only contains [ROOT histograms](https://root.cern.ch/root/htmldoc/guides/users-guide/Histograms.html) of the kinematic distributions, but also two [TTrees](https://root.cern.ch/doc/master/classTTree.html) of four-momentum tuples: one for data and one for fit intensities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.data import io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_data = io.pawian.read_hists_file(\"jpsi_f0_gammapipi.root\", \"data\")\n",
    "frame_fit = io.pawian.read_hists_file(\"jpsi_f0_gammapipi.root\", \"fit\")\n",
    "frame_fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how, here too, the {class}`~pandas.DataFrame` is formatted in such a way that it can be handled by the {class}`.PwaAccessor`. Also note that the {class}`~pandas.DataFrame` for the fit result contains weights. As discussed in {doc}`/usage/workflow/step4`, these are the fit intensities for each data point in the phase space. This allows us to already make some quick visualization of invariant mass distribution of the resonance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f0_data = frame_data[\"pi0_1\"] + frame_data[\"pi0_2\"]\n",
    "f0_fit = frame_fit[\"pi0_1\"] + frame_fit[\"pi0_2\"]\n",
    "\n",
    "f0_data.pwa.mass.hist(label=\"data\", bins=60, density=True, alpha=0.5)\n",
    "f0_fit.pwa.mass.hist(\n",
    "    label=\"fit\",\n",
    "    bins=60,\n",
    "    density=True,\n",
    "    histtype=\"step\",\n",
    "    color=\"red\",\n",
    "    weights=frame_fit.pwa.intensities,\n",
    ")\n",
    "\n",
    "plt.xlabel(\"$M(\\pi^0,\\pi^0)$ [GeV]\")\n",
    "plt.legend()\n",
    "plt.gca().set_title(\"f0 resonance\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also easily export the data again after you've made some adjustments, like selecting certain events. Just to illustrate the benefits of {mod}`pandas`, we apply some filter on one of the $\\pi^0$ mass, export the frame to an ASCII file and import it again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selection = frame_data[abs(f0_data.pwa.mass - 0.990) < 0.05]\n",
    "io.pawian.write_ascii(selection, \"filtered_data.dat\")\n",
    "io.pawian.read_ascii(\"filtered_data.dat\", [\"gamma\", \"pi0\", \"pi0\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conversion to kinematic variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having a {class}`PWA DataFrame <.PwaAccessor>`, we can use ComPWA to convert the momentum tuples to kinematic variables. For that, of course, we first need to {doc}`create a model file </usage/workflow/step1>` for the kinematics. As can be seen from the column names of the {class}`PWA DataFrame <.PwaAccessor>` that we imported from the `pawianHists.root` file, we have momentum tuples for a $J/\\psi \\to \\gamma\\pi^0\\pi^0$ decay and we saw that there is only one resonance ($f_0(980)$):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "import pycompwa.ui as pwa\n",
    "\n",
    "# not interested in warnings now\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.ERROR)\n",
    "pwa.Logging(\"ERROR\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.ui.system_control import (\n",
    "    InteractionTypes,\n",
    "    StateTransitionManager,\n",
    ")\n",
    "\n",
    "initial_state = [(\"J/psi\", [-1, 1])]\n",
    "final_state = [(\"gamma\"), (\"pi0\"), (\"pi0\")]\n",
    "tbd_manager = StateTransitionManager(\n",
    "    initial_state,\n",
    "    final_state,\n",
    "    formalism_type=\"helicity\",\n",
    "    topology_building=\"isobar\",\n",
    ")\n",
    "\n",
    "tbd_manager.set_allowed_interaction_types([InteractionTypes.EM])\n",
    "tbd_manager.allowed_intermediate_particles = [\"f0(980)\"]\n",
    "\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "solutions, _ = tbd_manager.find_solutions(graph_interaction_settings_groups)\n",
    "\n",
    "from pycompwa.expertsystem.amplitude.helicitydecay import (\n",
    "    HelicityAmplitudeGeneratorXML,\n",
    ")\n",
    "\n",
    "model_file = \"jpsi_f0_gammapipi.xml\"\n",
    "xml_generator = HelicityAmplitudeGeneratorXML()\n",
    "xml_generator.generate(solutions)\n",
    "xml_generator.write_to_file(model_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have an XML model file defining the kinematics and a {class}`PWA DataFrame <.PwaAccessor>`, we can use the convert module to convert the {class}`~pandas.DataFrame` to an {class}`.EventCollection`. Note, however, that we will run into an exception:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.data import convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    convert.pandas_to_events(frame_data, model_file)\n",
    "except Exception as exc:\n",
    "    print(\"EXCEPTION:\", exc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's going on here? The kinematics file works with final state IDs, so it doesn't understand the particle names here. Now, we could try to follow the first suggestion here, but this won't work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.data import naming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "naming.particle_to_id(frame_data, model_file)\n",
    "frame_data.pwa.particles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the $\\gamma$ has been nicely renamed to its final state ID, but the renaming failed for the pions (it would have worked if the separator used for the added index for duplicate particles were a `-`). If we follow the second suggestion, it will work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mapping = {\"gamma\": 2, \"pi0_1\": 3, \"pi0_2\": 4}\n",
    "frame_data.rename(columns=mapping, inplace=True)\n",
    "frame_fit.rename(columns=mapping, inplace=True)\n",
    "events_data = convert.pandas_to_events(frame_data, model_file)\n",
    "events_fit = convert.pandas_to_events(frame_fit, model_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have {class}`.EventCollection` instances, you are free to use all ComPWA functionality from {doc}`/usage/workflow/step3` onwards. If, however, you were more interested in the kinematic variables for these imported data sets immediately, you can expand the original {class}`PWA DataFrame <.PwaAccessor>` with the kinematic variables as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycompwa.ui as pwa\n",
    "from pycompwa.data import append"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_data = pwa.compute_kinematic_variables(events_data, model_file)\n",
    "set_fit = pwa.compute_kinematic_variables(events_fit, model_file)\n",
    "naming.id_to_particle(frame_data, model_file, make_unique=True)\n",
    "naming.id_to_particle(frame_fit, model_file, make_unique=True)\n",
    "append(frame_data, convert.data_set_to_pandas(set_data))\n",
    "append(frame_fit, convert.data_set_to_pandas(set_fit))\n",
    "frame_data.pwa.other_columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can plot the distributions of the kinematic variables (as computed by ComPWA) **of the imported data**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = \"theta_2_4_vs_3\"\n",
    "frame_data[var].hist(label=\"data\", bins=60, density=True, alpha=0.5)\n",
    "frame_fit[var].hist(\n",
    "    label=\"fit\",\n",
    "    bins=60,\n",
    "    density=True,\n",
    "    histtype=\"step\",\n",
    "    color=\"red\",\n",
    "    weights=frame_fit.pwa.intensities,\n",
    ")\n",
    "plt.gca().set_title(naming.replace_ids(var, model_file))\n",
    "plt.legend();"
   ]
  }
 ],
 "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
}
