{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Create an amplitude model"
   ]
  },
  {
   "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": [
    "Your analysis starts by defining an amplitude model that describes the reaction process you want to study. Such a model is generally very complex and requires a fair amount of effort by the analyst (you). This also gives a lot of room for mistakes.\n",
    "\n",
    "To facilitate this process of defining an amplitude model, {mod}`pycompwa` comes with a so-called [‘expert system’](https://en.wikipedia.org/wiki/Expert_system). Its responsibility is to give you advice on the form of an amplitude model based on the problem set you define (initial state, final state, allowed interactions, intermediate states, etc.). Internally, the system propagates the quantum numbers through the reaction graph while satisfying the specified conservation rules. How to control this procedure is explained in more detail below.\n",
    "\n",
    "Afterwards, the amplitude model of the expert system can be exported. The model can for instance be used to generate a data set (toy Monte Carlo) for this reaction and to optimize its parameters to resemble an actual data set as good as possible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1. Define the problem set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first define the boundary conditions of our physics problem, such as initial state, final state, formalism type, etc. and pass all of that information to the {class}`.StateTransitionManager`. This is the main user interface class of the ComPWA expert system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.ui.system_control import StateTransitionManager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_state = [(\"J/psi\", [-1, 1])]\n",
    "final_state = [(\"gamma\", [-1, 1]), (\"pi0\", [0]), (\"pi0\", [0])]\n",
    "\n",
    "tbd_manager = StateTransitionManager(\n",
    "    initial_state,\n",
    "    final_state,\n",
    "    formalism_type=\"helicity\",\n",
    "    topology_building=\"isobar\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "The {class}`.StateTransitionManager` (STM) is the main user interface class of the ComPWA expert system. The boundary conditions of your physics problem, such as the initial state, final state, formalism type, etc., are defined through this interface.\n",
    "\n",
    "* {meth}`.prepare_graphs` of the STM creates all topology graphs ― here, using the isobar model (a tree of two-body decays). The function also initializes the graphs with the initial and final state and a set of conservation laws at each interaction node.\n",
    "\n",
    "* By default, all three interaction types (strong, EM, weak) are used in the preparation stage. However, it is also possible to choose the allowed interaction types globally via {meth}`.set_allowed_interaction_types`.\n",
    "\n",
    "After the preparation step, you can modify the settings returned by {meth}`.prepare_graphs` to your liking. Since the output of this function contains quite a lot of information, the {mod}`expertsystem UI <pycompwa.expertsystem.ui>` aids in the configuration (especially the STM).\n",
    "\n",
    "* A subset of particles that are allowed as intermediate states can also be specified: either through the {class}`STM's constructor <.StateTransitionManager>` or by setting the instance attribute `allowed_intermediate_particles`.\n",
    "\n",
    ":::{hint}\n",
    "\n",
    "How does the {class}`.StateTransitionManager` know what a `\"J/psi\"` is? It simply uses a default particle list that is shipped with ComPWA (see [here](https://github.com/ComPWA/ComPWA/blob/master/Physics/particle_list.xml)) and is handled by the {mod}`.particle` module. For more info, see {doc}`../tools/particles`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2. Prepare topologies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create all topology graphs using the **isobar model** (tree of two-body decays) and initialize the graphs with the initial and final state. Remember that each interaction node defines its own set of conservation laws.\n",
    "\n",
    "The {class}`.StateTransitionManager` (STM) defines three interaction types:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Interaction          | Strength  |\n",
    "| -------------------- | --------- |\n",
    "| strong               | $60$      |\n",
    "| electromagnetic (EM) | $1$       |\n",
    "| weak                 | $10^{-4}$ |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, all three are used in the preparation stage. The function {meth}`.prepare_graphs` of the STM generates graphs with all possible combinations of interaction nodes. An overall interaction strength is assigned to each graph and they are grouped according to this strength."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3. Find solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are happy with the default settings generated by the {class}`.StateTransitionManager`, just start with solving directly!\n",
    "\n",
    ":::{note}\n",
    "\n",
    "This step takes about 30 sec on an Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz running multi-threaded.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solutions, violated_rules = tbd_manager.find_solutions(\n",
    "    graph_interaction_settings_groups\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.topology.graph import get_intermediate_state_edges\n",
    "\n",
    "\n",
    "def print_intermediate_states(solutions):\n",
    "    \"\"\"Just a little function to print the intermediate states.\"\"\"\n",
    "    print(\"intermediate states:\")\n",
    "    intermediate_states = set()\n",
    "    for g in solutions:\n",
    "        edge_id = get_intermediate_state_edges(g)[0]\n",
    "        intermediate_states.add(g.edge_props[edge_id][\"@Name\"])\n",
    "    print(intermediate_states)\n",
    "\n",
    "\n",
    "print(\"found\", len(solutions), \"solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a lot of solutions that are actually heavily suppressed (they involve two weak decays).\n",
    "\n",
    "In general, you can modify the dictionary returned by {meth}`.prepare_graphs` directly, but the STM also comes with functionality to globally choose the allowed interaction types.\n",
    "\n",
    "So, go ahead and **disable** the **EM** and **weak** interaction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.ui.system_control import InteractionTypes\n",
    "\n",
    "tbd_manager.set_allowed_interaction_types([InteractionTypes.Strong])\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "solutions, violated_rules = tbd_manager.find_solutions(\n",
    "    graph_interaction_settings_groups\n",
    ")\n",
    "\n",
    "print(\"found\", len(solutions), \"solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Huh, what happened here? Actually, since a $\\gamma$ particle appears in one of the interaction nodes, the expert system knows that this node **must involve EM interactions**! Because the node can be an effective interaction, the weak interaction cannot be excluded, as it contains only a subset of conservation laws.\n",
    "\n",
    "Since only the strong interaction was supposed to be used, this results in a warning and the STM automatically corrects the mistake.\n",
    "\n",
    "Once the EM interaction is included, this warning disappears. Be aware, however, that the EM interaction is now available globally. Hence, there now might be solutions in which both nodes are electromagnetic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tbd_manager.set_allowed_interaction_types(\n",
    "    [InteractionTypes.Strong, InteractionTypes.EM]\n",
    ")\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "solutions, violated_rules = tbd_manager.find_solutions(\n",
    "    graph_interaction_settings_groups\n",
    ")\n",
    "\n",
    "print(\"found\", len(solutions), \"solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Now we selected only the strongest contributions. Be aware, though, that there are more effects that can suppress certain decays, like small branching ratios. In this example, the initial state $J/\\Psi$ can decay into $\\pi^0 + \\rho^0$ or $\\pi^0 + \\omega$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| decay                               | branching ratio |\n",
    "| ----------------------------------- | --------------- |\n",
    "| $\\omega \\rightarrow \\gamma+\\pi^0$ | 0.0828          |\n",
    "| $\\rho^0 \\rightarrow \\gamma+\\pi^0$ | 0.0006          |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, the $\\rho^0$ mainly decays into $\\pi+\\pi$, not $\\gamma+\\pi^0$ and is therefore suppressed. This information is currently not known to the expert system, but it is possible to hand the expert system a list of allowed intermediate states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# particles are found by name comparison,\n",
    "# i.e. f2 will find all f2's and f all f's independent of their spin\n",
    "tbd_manager.allowed_intermediate_particles = [\"f\"]\n",
    "\n",
    "solutions, violated_rules = tbd_manager.find_solutions(\n",
    "    graph_interaction_settings_groups\n",
    ")\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have selected all amplitudes that involve **f** states. The warnings appear only to notify the user that the list of solutions is not exhaustive: for certain edges in the graph, no suitable particle was found (since only f states were allowed).\n",
    "\n",
    "Now that we are satisfied with the intermediate resonances, we are all set to generate an amplitude model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.amplitude.helicitydecay import (\n",
    "    HelicityAmplitudeGeneratorXML,\n",
    ")\n",
    "\n",
    "xml_generator = HelicityAmplitudeGeneratorXML()\n",
    "xml_generator.generate(solutions)\n",
    "xml_generator.write_to_file(\"model.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "In this example, we used the helicity formalism. If you want to work with the canonical formalism, you have to:\n",
    "\n",
    "1. Construct a {class}`.StateTransitionManager` with argument `formalism_type='canonical-helicity'` instead of `formalism_type='helicity'`.\n",
    "2. Use the {class}`.CanonicalAmplitudeGeneratorXML` instead of the {class}`.HelicityAmplitudeGeneratorXML`.\n",
    "\n",
    "Note that you **have to do both**, because otherwise the resulting amplitude will not make sense!\n",
    "\n",
    ":::\n",
    "\n",
    "Have a look through the sections of the resulting XML file to see what you recognize from the problem set defined above. There may also be some things you want to change in there manually, so **make sure you store this XML file** carefully (e.g. track it with Git) so that you don't overwrite it your changed after rerunning the expert system.\n",
    "\n",
    "In the next step, we'll use this amplitude model to generate a data 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
}
