{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell",
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# WARNING: advised to install a specific version, e.g. tensorwaves==0.1.2\n",
    "%pip install -q tensorwaves[doc,jax,viz] IPython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "import os\n",
    "\n",
    "from IPython.display import display  # noqa: F401\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 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",
    "from ampform.helicity import HelicityModel\n",
    "\n",
    "with open(\"helicity_model.pickle\", \"rb\") as model_file:\n",
    "    model: HelicityModel = 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}`~ampform.kinematics.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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from tensorwaves.data import generate_phsp\n",
    "\n",
    "phsp_sample = generate_phsp(100_000, model.adapter.reaction_info)\n",
    "pd.DataFrame(phsp_sample.to_pandas())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "The resulting phase space sample is a {class}`~ampform.data.EventCollection` of {class}`~ampform.data.FourMomentumSequence`s for each particle in the final state. The {meth}`~ampform.data.EventCollection.to_pandas` method can be used to cast the {class}`~ampform.data.EventCollection` to a format that can be understood by {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": {},
   "outputs": [],
   "source": [
    "from tensorwaves.model import LambdifiedFunction, SympyModel  # noqa: F401\n",
    "\n",
    "sympy_model = SympyModel(\n",
    "    expression=model.expression,\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": {},
   "outputs": [],
   "source": [
    "from tensorwaves.data import generate_data\n",
    "\n",
    "data_sample = generate_data(\n",
    "    size=10_000,\n",
    "    reaction_info=model.adapter.reaction_info,\n",
    "    data_transformer=data_converter,\n",
    "    intensity=intensity,\n",
    ")\n",
    "pd.DataFrame(data_sample.to_pandas())"
   ]
  },
  {
   "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 state_id, particle 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": [
    "Just like {obj}`~ampform.data.EventCollection`, the {obj}`~ampform.data.DataSet` can easily be converted it to a {class}`pandas.DataFrame`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "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": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data_frame[\"m_12\"].hist(bins=100, alpha=0.5, density=True, figsize=(8, 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",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}