{
 "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"
    ]
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%config Completer.use_jedi = False\n",
    "%config InlineBackend.figure_formats = ['svg']\n",
    "import os\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)\n",
    "\n",
    "# Install on Google Colab\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "from IPython import get_ipython\n",
    "\n",
    "install_packages = \"google.colab\" in str(get_ipython())\n",
    "if install_packages:\n",
    "    for package in [\"tensorwaves[doc]\", \"graphviz\"]:\n",
    "        subprocess.check_call(\n",
    "            [sys.executable, \"-m\", \"pip\", \"install\", package]\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Generate data samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will use the {class}`~ampform.helicity.HelicityModel` that we created with {mod}`ampform` in [the previous step](step1) to generate a data sample via hit & miss Monte Carlo. We do this with the {mod}`.data` module.\n",
    "\n",
    "First, we {func}`~pickle.load` the {class}`~ampform.helicity.HelicityModel` that was created in the previous step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"helicity_model.pickle\", \"rb\") as model_file:\n",
    "    model = pickle.load(model_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "reaction_info = model.adapter.reaction_info\n",
    "initial_state = next(iter(reaction_info.initial_state.values()))\n",
    "print(\"Initial state:\")\n",
    "print(\" \", initial_state.name)\n",
    "print(\"Final state:\")\n",
    "for i, p in reaction_info.final_state.items():\n",
    "    print(f\"  {i}: {p.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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 the {func}`.generate_phsp` function. By default, this function uses {class}`.TFPhaseSpaceGenerator` as a, well... phase-space generator (using {obj}`tensorflow <tf.Tensor>` and the [`phasespace`](https://phasespace.readthedocs.io) package as a back-end) and generates random numbers with {class}`.TFUniformRealNumberGenerator`. You can use other generators with the arguments of {func}`.generate_phsp`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As opposed to the main {ref}`usage:Generate data sample` of the main usage example page, we will generate a **deterministic** data sample. This can be done by feeding a {class}`.UniformRealNumberGenerator` with a specific {attr}`.UniformRealNumberGenerator.seed` and feeding that generator to the {func}`.generate_phsp` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from tensorwaves.data import TFUniformRealNumberGenerator, generate_phsp\n",
    "\n",
    "rng = TFUniformRealNumberGenerator(seed=0)\n",
    "reaction_info = model.adapter.reaction_info\n",
    "initial_state_mass = reaction_info.initial_state[-1].mass\n",
    "final_state_masses = {i: p.mass for i, p in reaction_info.final_state.items()}\n",
    "phsp_sample = generate_phsp(\n",
    "    size=100_000,\n",
    "    initial_state_mass=initial_state_mass,\n",
    "    final_state_masses=final_state_masses,\n",
    "    random_generator=rng,\n",
    ")\n",
    "pd.DataFrame(\n",
    "    {\n",
    "        (k, label): np.transpose(v)[i]\n",
    "        for k, v in phsp_sample.items()\n",
    "        for i, label in enumerate([\"E\", \"px\", \"py\", \"pz\"])\n",
    "    }\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`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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 with the function {func}`.generate_data`. Its usage is similar to {func}`.generate_phsp`, 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 {doc}`step1`, 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 an {class}`.Model` object that it can then {meth}`~.Model.lambdify` to a {obj}`.Function` object.\n",
    "\n",
    "The {obj}`~ampform.helicity.HelicityModel`  was expressed in terms of {mod}`sympy`, so we express the model as a {class}`.SympyModel` and lambdify it to a {class}`.LambdifiedFunction`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "\n",
    "Here, we make use of {func}`.optimized_lambdify` by specifying `max_complexity`. See {ref}`usage/faster-lambdify:Specifying complexity`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip-flake8"
    ]
   },
   "outputs": [],
   "source": [
    "from tensorwaves.model import LambdifiedFunction, SympyModel\n",
    "\n",
    "sympy_model = SympyModel(\n",
    "    expression=model.expression.doit(),\n",
    "    parameters=model.parameter_defaults,\n",
    "    max_complexity=200,\n",
    ")\n",
    "%time intensity = LambdifiedFunction(sympy_model, backend=\"numpy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A problem is that {class}`.LambdifiedFunction` takes a {obj}`.DataSample` 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}`.HelicityTransformer`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorwaves.data.transform import HelicityTransformer\n",
    "\n",
    "data_converter = HelicityTransformer(model.adapter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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 {ref}`phase-space sample we generated previously <usage/step2:2.1 Generate phase space sample>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "from tensorwaves.data import generate_data\n",
    "\n",
    "data_sample = generate_data(\n",
    "    size=10_000,\n",
    "    initial_state_mass=initial_state_mass,\n",
    "    final_state_masses=final_state_masses,\n",
    "    data_transformer=data_converter,\n",
    "    intensity=intensity,\n",
    "    random_generator=rng,\n",
    ")\n",
    "pd.DataFrame(\n",
    "    {\n",
    "        (k, label): np.transpose(v)[i]\n",
    "        for k, v in data_sample.items()\n",
    "        for i, label in enumerate([\"E\", \"px\", \"py\", \"pz\"])\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we use a {class}`.UniformRealNumberGenerator` with a specific {attr}`.UniformRealNumberGenerator.seed` to ensure we get a **deterministic** data sample."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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}`.HelicityTransformer` 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_set = data_converter.transform(phsp_sample)\n",
    "data_set = data_converter.transform(data_sample)\n",
    "list(data_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {obj}`~ampform.data.DataSet` is just a mapping of kinematic variables names to a sequence 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 (\n",
    "    state_id,\n",
    "    particle,\n",
    ") in model.adapter.reaction_info.final_state.items():\n",
    "    print(f\"ID {state_id}:\", particle.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "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",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {obj}`~ampform.data.DataSet` can easily be converted it 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_set)\n",
    "phsp_frame = pd.DataFrame(data_set)\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": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data_frame[\"m_12\"].hist(bins=100, alpha=0.5, density=True, figsize=(9, 4))\n",
    "plt.xlabel(\"$m$ [GeV]\")\n",
    "for i, p in enumerate(intermediate_states):\n",
    "    plt.axvline(x=p.mass, linestyle=\"dotted\", label=p.name, color=colors[i])\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{ref}`usage/step3:Intensity components`\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 Export data sets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorWaves currently has no export functionality for data samples, so we just {func}`pickle.dump` these data samples as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"data_set.pickle\", \"wb\") as stream:\n",
    "    pickle.dump(data_set, stream)\n",
    "with open(\"phsp_set.pickle\", \"wb\") as stream:\n",
    "    pickle.dump(phsp_set, stream)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the {doc}`next step <step3>`, we illustrate how to {meth}`~.Minuit2.optimize` the intensity model to these data samples."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}