{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate data samples"
   ]
  },
  {
   "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": [
    "There are two different types of data samples: (1) {ref}`phase space samples <usage/tools/generate:Phase space samples>` and (2) {ref}`intensity-based samples <usage/tools/generate:Intensity-based samples>` (which mimic the intensity distributions of the reaction). In addition, you may want to generate more data points in specific regions of phase space (e.g. in the case of a sharp resonance). This is called {ref}`importance sampling <usage/tools/generate:Importance sampling>`.\n",
    "\n",
    "Usually, you need a phase space sample to generate an intensity-based sample, so we start with that."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Phase space samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the {mod}`pycompwa.ui` module you can easily create various types of Monte Carlo phase space samples. You do this with the function {func}`.generate_phsp`.\n",
    "\n",
    "One of the required arguments is a {class}`.PhaseSpaceEventGenerator` instance. There are currently two options: {class}`.EvtGenGenerator` and {class}`.RootGenerator`. Using the {class}`.EvtGenGenerator` is recommended for numerical precision:\n",
    "\n",
    "```{eval-rst}\n",
    ".. autoclass:: pycompwa.ui.EvtGenGenerator\n",
    "    :special-members: __init__\n",
    "    :noindex:\n",
    "```\n",
    "\n",
    "As you can see, such a {class}`.PhaseSpaceEventGenerator` requires state transition info from a {class}`.Kinematics` instance, such as the {class}`.HelicityKinematics` class. \n",
    "\n",
    "The second argument for the {func}`.generate_phsp` function is a random number generator. Here, there are again two options: {class}`.StdUniformRealGenerator` and {class}`.RootUniformRealGenerator`.\n",
    "\n",
    "To illustrate this procedure, we first need to {doc}`create some kinematics model </usage/workflow/step1>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.ERROR)  # not interested in warnings now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.amplitude.helicitydecay import (\n",
    "    HelicityAmplitudeGeneratorXML,\n",
    ")\n",
    "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(\n",
    "    [InteractionTypes.Strong, InteractionTypes.EM]\n",
    ")\n",
    "tbd_manager.allowed_intermediate_particles = [\"f\"]\n",
    "\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "solutions, _ = tbd_manager.find_solutions(graph_interaction_settings_groups)\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "\n",
    "xml_generator = HelicityAmplitudeGeneratorXML()\n",
    "xml_generator.generate(solutions)\n",
    "xml_generator.write_to_file(\"model.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can go through the procedure explained above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycompwa.ui as pwa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify kinematics\n",
    "kinematics = pwa.create_helicity_kinematics(\"model.xml\")\n",
    "\n",
    "# specify generators\n",
    "generator = pwa.EvtGenGenerator(\n",
    "    kinematics.get_particle_state_transition_kinematics_info()\n",
    ")\n",
    "random_generator = pwa.StdUniformRealGenerator(12345)\n",
    "\n",
    "# generate phase space sample\n",
    "phsp_sample = pwa.generate_phsp(1000, generator, random_generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intensity-based samples"
   ]
  },
  {
   "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 an {class}`.Intensity` object and a phase space over which to generate that intensity distribution. We therefore call such a data sample an **intensity-based sample**.\n",
    "\n",
    "An intensity-based sample is generated with the function {func}`~.ui.generate`. Its usage is similar to {func}`.generate_phsp`, but now you have to give an {class}`.Intensity` and a {class}`.Kinematics` instance as additional arguments. Sampling is performed through hit & miss, which is reflected by the efficiency printed at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create intensity profile\n",
    "particle_list = pwa.read_particles(\"model.xml\")\n",
    "intensity_builder = pwa.IntensityBuilderXML(\n",
    "    \"model.xml\", particle_list, kinematics, phsp_sample\n",
    ")\n",
    "intensity = intensity_builder.create_intensity()\n",
    "\n",
    "# generate data\n",
    "data_sample = pwa.generate(\n",
    "    100, kinematics, generator, intensity, random_generator\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importance sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, we used the decay $J/\\psi \\to \\gamma \\pi^0 \\pi^0$ with f resonances only. Imagine, however, that we want to investigate the resonance $\\omega \\to \\gamma\\pi^0$. This resonance has such a narrow peak that you will run into problems when normalizing the intensity model.\n",
    "\n",
    "To normalize a distribution, you need know the integral over that distribution. The integral over an intensity distribution distribution can only be computed numerically. We simply do this by computing the sum of the intensities evaluated for each point in a phase space (Monte Carlo integration), and averaging over the number of phase space points. This means that the higher the number of phase space points, the more accurate your integral.\n",
    "\n",
    "With a narrow peak, however, many of the points in an evenly distributed phase space sample will lie outside the peak region, where the intensity (and therefore their contribution to the integral) is highest. So in order to have a good estimate for the integral over this narrow peak structure, you want to have a larger number of phase space points under that peak. Now, you could just use a bigger phase space sample (increasing number of points within the peak region), but this will be resource intensive as it will result in many hits in regions that hardly contributes to the integral.\n",
    "\n",
    "A smarter strategy is generate a sample that has a higher density of phase space points within the peak region. This strategy is called **importance sampling**.\n",
    "\n",
    "We first again create an amplitude model, this time including a 'widened' $\\omega$ resonance (see also {doc}`particles`). Widening $\\omega$ is only done so that the resonance is still visible―the peak is still extremely narrow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.amplitude.helicitydecay import (\n",
    "    HelicityAmplitudeGeneratorXML,\n",
    ")\n",
    "from pycompwa.expertsystem.state.particle import particle_list\n",
    "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",
    "# Widen omega resonance\n",
    "omega = particle_list[\"omega(782)\"]\n",
    "parameters = omega[\"DecayInfo\"][\"Parameter\"]\n",
    "for par in parameters:\n",
    "    if par[\"@Type\"] == \"Width\":\n",
    "        par[\"Value\"] = 0.001\n",
    "\n",
    "tbd_manager.set_allowed_interaction_types(\n",
    "    [InteractionTypes.Strong, InteractionTypes.EM]\n",
    ")\n",
    "tbd_manager.allowed_intermediate_particles = [\"f2(1270)\", \"omega\"]\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "solutions, _ = tbd_manager.find_solutions(graph_interaction_settings_groups)\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "\n",
    "xml_generator = HelicityAmplitudeGeneratorXML()\n",
    "xml_generator.generate(solutions)\n",
    "xml_generator.write_to_file(\"model.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we generate a phase space sample that is importance sampled by the\n",
    "intensity. This can be done with the\n",
    "{func}`.generate_importance_sampled_phsp` function.\n",
    "\n",
    ":::{note}\n",
    "\n",
    "Note that you still need a 'normal' phase space sample for the {class}`.IntensityBuilderXML`. In theory, however, this shouldn't be necessary at this stage, but in the current implementation, this sample is required for computing the outermost intensity. Luckily, during data generation the parameters of the model stay fixed, so this normalization is irrelevant and just causes a constant scaling factor.\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": [
    "# specify kinematics\n",
    "particle_list = pwa.read_particles(\"model.xml\")\n",
    "kinematics = pwa.create_helicity_kinematics(\"model.xml\", particle_list)\n",
    "\n",
    "# specify generators\n",
    "generator = pwa.EvtGenGenerator(\n",
    "    kinematics.get_particle_state_transition_kinematics_info()\n",
    ")\n",
    "random_generator = pwa.StdUniformRealGenerator(12345)\n",
    "\n",
    "# generate evenly distributed phase space sample\n",
    "phsp_sample = pwa.generate_phsp(100000, generator, random_generator)\n",
    "\n",
    "# create intensity\n",
    "intensity_builder = pwa.IntensityBuilderXML(\n",
    "    \"model.xml\", particle_list, kinematics, phsp_sample\n",
    ")\n",
    "intensity = intensity_builder.create_intensity()\n",
    "\n",
    "# generate importance-sampled phase space sample\n",
    "phsp_sample_importance = pwa.generate_importance_sampled_phsp(\n",
    "    2000, kinematics, generator, intensity, random_generator\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "\n",
    "The current implementation is simple and uniformly generates events in the phase space domain. Depending on the shape of your Intensity, this can be quite computation intensive. However, ideally you would only run this once for a specific reaction.\n",
    "\n",
    ":::\n",
    "\n",
    "The effect of this importance sampling can be best seen in a Dalitz plot. There, one can see that there are more events in the regions where the intensity is large (the resonances), but one still sees a uniform distribution overall.\n",
    "\n",
    "To create a Dalitz plot, we first have to create a {class}`.DataSet`, as we saw in {ref}`this section <usage/workflow/step4:4.2 Visualize>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kinematics.create_all_subsystems()\n",
    "phsp_set = kinematics.convert(phsp_sample)\n",
    "phsp_set_importance = kinematics.convert(phsp_sample_importance)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use {mod}`matplotlib.pyplot` to create the Dalitz plot. Note that we use {func}`.replace_ids` to transform the final state IDs into readable particles names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "from pycompwa.data import naming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dalitz_plot(data_set, mass_x, mass_y, bins=50, **kwargs):\n",
    "    \"\"\"Helper function to create a Dalitz plot with useful axis titles.\"\"\"\n",
    "    plt.hist2d(\n",
    "        data_set.data[mass_x],\n",
    "        data_set.data[mass_y],\n",
    "        weights=data_set.weights,\n",
    "        bins=bins,\n",
    "        **kwargs,\n",
    "    )\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": [
    "dalitz_plot(\n",
    "    phsp_set_importance, \"mSq_(3,4)\", \"mSq_(2,4)\", bins=80, norm=LogNorm()\n",
    ")\n",
    "plt.gca().set_title(\"importance-sampled phase space\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how this compares to the evenly distributed phase space sample and the intensity-based sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dalitz_plot(phsp_set, \"mSq_(3,4)\", \"mSq_(2,4)\", bins=80, norm=LogNorm())\n",
    "plt.gca().set_title(\"evenly distributed phase space\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwa.Logging(\"error\")\n",
    "data_sample = pwa.generate(\n",
    "    2000, kinematics, generator, intensity, random_generator\n",
    ")\n",
    "data_set = kinematics.convert(data_sample)\n",
    "dalitz_plot(data_set, \"mSq_(3,4)\", \"mSq_(2,4)\", bins=80, norm=LogNorm())\n",
    "plt.gca().set_title(\"intensity-based sample\");"
   ]
  }
 ],
 "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
}
