{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "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",
    "\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 [\"expertsystem\", \"graphviz\"]:\n",
    "        subprocess.check_call(\n",
    "            [sys.executable, \"-m\", \"pip\", \"install\", package]\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Formulate amplitude model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{warning}\n",
    "The {doc}`PWA Expert System <index>` has been split up into {doc}`QRules <qrules:index>` and {doc}`AmpForm <ampform:index>`. Please use these packages instead!\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate transitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In {doc}`reaction`, we use {func}`.generate_transitions` to create a list of allowed {class}`.StateTransitionGraph` instances for a specific decay channel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import expertsystem as es\n",
    "\n",
    "result = es.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [-1, +1]),\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\"f(0)(980)\", \"f(0)(1500)\"],\n",
    "    allowed_interaction_types=[\"strong\", \"EM\"],\n",
    "    formalism_type=\"canonical-helicity\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import graphviz\n",
    "\n",
    "dot = es.io.asdot(result, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build SymPy expression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use the {class}`.Result` to formulate an amplitude model. The type of this amplitude model is dependent on the {attr}`~.Result.formalism_type`. The function {func}`.amplitude.get_builder` helps to get the correct amplitude builder class for this {attr}`~.Result.formalism_type`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_builder = es.amplitude.get_builder(result)\n",
    "type(model_builder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we now use the {meth}`.HelicityAmplitudeBuilder.generate` method of this builder, we get an amplitude model without any dynamics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_no_dynamics = model_builder.generate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {attr}`.HelicityModel.expression` is just a {class}`sympy.Expr <sympy.core.expr.Expr>`, which we can pull apart by using its {attr}`~sympy.core.basic.Basic.args` (see {doc}`sympy:tutorial/manipulation`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensities = model_no_dynamics.expression.args\n",
    "intensity_1 = intensities[0]\n",
    "base, power = intensity_1.args\n",
    "abs_arg = base.args[0]\n",
    "amplitudes = abs_arg.args\n",
    "amplitudes[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To set dynamics for specific resonances, use {meth}`.set_dynamics` on the same {class}`.HelicityAmplitudeBuilder` instance.  You can set the dynamics to be any kind of {class}`~sympy.core.expr.Expr`, as long as you keep track of which {class}`~sympy.core.symbol.Symbol` names you use. The {mod}`expertsystem` does provide a few common {mod}`.lineshape` functions however, which can be constructed as {class}`~sympy.core.expr.Expr` with the correct {class}`~sympy.core.symbol.Symbol` names using {meth}`.set_dynamics`. This function takes specific {mod}`.builder` functions, such as {func}`.create_relativistic_breit_wigner`, which would create a {func}`.relativistic_breit_wigner` for a specific resonance. Here's an example for a relativistic Breit-Wigner _with form factor_ for the intermediate resonances and use a Blatt-Weisskopf barrier factor for the production decay:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.amplitude.dynamics.builder import (\n",
    "    create_non_dynamic_with_ff,\n",
    "    create_relativistic_breit_wigner_with_ff,\n",
    ")\n",
    "\n",
    "initial_state_particle = result.get_initial_state()[0]\n",
    "model_builder.set_dynamics(initial_state_particle, create_non_dynamic_with_ff)\n",
    "for name in result.get_intermediate_particles().names:\n",
    "    model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{seealso}\n",
    "{doc}`dynamics/custom`\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we use the reconfigured {class}`.HelicityAmplitudeBuilder` to generate another {class}`.HelicityModel`, this time with relativistic Breit-Wigner functions and form factors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model_builder.generate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mathematical formula"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's possible to view the complete amplitude model as an expression. This would, however, clog the screen here, so we instead just view the formula of one of its {attr}`~.HelicityModel.components`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "some_amplitude = list(model.components.values())[0]\n",
    "some_amplitude.doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "We use {meth}`~sympy.core.basic.Basic.doit` to evaluate the Wigner-$D$ ({meth}`Rotation.D <sympy.physics.quantum.spin.Rotation.D>`) and Clebsch-Gordan ({class}`~sympy.physics.quantum.cg.CG`) functions in the full expression.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {attr}`.HelicityModel.parameter_defaults` attribute can be used to substitute all parameters with suggested values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "some_amplitude.doit().subs(model.parameter_defaults)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case ($J/\\psi \\to \\gamma f_0, f_0 \\to \\pi^0\\pi^0$) _without dynamics_, the total intensity is only dependent on the $\\theta$ angle of $\\gamma$ in the center of mass frame (see {func}`.get_helicity_angle_label`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_dynamics = model_no_dynamics.expression.doit()\n",
    "no_dynamics_substituted = no_dynamics.subs(model.parameter_defaults)\n",
    "no_dynamics_substituted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(no_dynamics_substituted.free_symbols) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import sympy as sy\n",
    "\n",
    "theta = next(iter(no_dynamics_substituted.free_symbols))\n",
    "sy.plot(\n",
    "    no_dynamics_substituted,\n",
    "    (theta, 0, sy.pi),\n",
    "    axis_center=(0, 0),\n",
    "    ylabel=\"$I$\",\n",
    "    ylim=(0, 16),\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this decay channel, the amplitude model is built up of four components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_dynamics.subs(zip(no_dynamics.args, sy.symbols(\"I_{:4}\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can be nicely visualized as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import sympy as sy\n",
    "\n",
    "plots = []\n",
    "colors = [\"red\", \"blue\", \"green\", \"purple\"]\n",
    "\n",
    "total = 0\n",
    "for i, intensity in enumerate(no_dynamics.args):\n",
    "    total += intensity.subs(model.parameter_defaults).doit()\n",
    "    plots.append(\n",
    "        sy.plot(\n",
    "            total,\n",
    "            (theta, 0, sy.pi),\n",
    "            axis_center=(0, 0),\n",
    "            ylabel=\"$I$\",\n",
    "            ylim=(0, 16),\n",
    "            line_color=colors[i],\n",
    "            show=False,\n",
    "            label=f\"$I_{i}$\",\n",
    "            legend=True,\n",
    "        )\n",
    "    )\n",
    "for i in range(1, 4):\n",
    "    plots[0].extend(plots[i])\n",
    "plots[0].show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the model _with dynamics_, we have several free symbols, such as the mass and width of the resonances. For the fitting package these will be considered **parameters** that are to be optimized and (kinematic) **variables** that represent the data set. Examples of parameters are mass ($m_\\text{particle}$) and width ($\\Gamma_\\text{particle}$) of the resonances and certain amplitude coefficients ($C$). Examples of kinematic variables are the helicity angles $\\theta$ and $\\phi$ and the invariant masses ($m_{ij...}$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(model.expression.free_symbols, key=lambda s: s.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say we want to plot the amplitude model with respect to $m_{3+4}$. We would have to substitute all other free symbols with some value."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we can use {attr}`.HelicityModel.parameter_defaults` to substitute the parameters with suggested values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "suggested_expression = model.expression.subs(model.parameter_defaults)\n",
    "suggested_expression.free_symbols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ideally, we would now 'integrate out' the helicity angles. Here, we however just set these angles to $0$, as computing the integral would take quite some time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "angle = 0\n",
    "angle_substitutions = {\n",
    "    s: angle\n",
    "    for s in suggested_expression.free_symbols\n",
    "    if s.name.startswith(\"phi\") or s.name.startswith(\"theta\")\n",
    "}\n",
    "evaluated_angle_intensity = suggested_expression.subs(angle_substitutions)\n",
    "evaluated_angle_intensity.free_symbols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By now we are only left with the masses of the final state particles ($m_1$ and $m_2$), since they appear as symbols in the {func}`.relativistic_breit_wigner_with_ff`. Final state particles 3 and 4 are the $\\pi^0$'s, so we can just substitute them with their masses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.reaction.particle import load_pdg\n",
    "\n",
    "pi0 = load_pdg()[\"pi0\"]\n",
    "plotted_intensity = evaluated_angle_intensity.doit().subs(\n",
    "    {\n",
    "        sy.Symbol(\"m_1\", real=True): pi0.mass,\n",
    "        sy.Symbol(\"m_2\", real=True): pi0.mass,\n",
    "    },\n",
    "    simultaneous=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{tip}\n",
    "Use {meth}`~sympy.core.basic.Basic.subs` with `simultaneous=True`, as that avoids a bug later on when plotting with {mod}`matplotlib.pyplot`.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! Now we are only left with the invariant mass $m_{3+4}$ of the two pions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(plotted_intensity.free_symbols) == 1\n",
    "m = next(iter(plotted_intensity.free_symbols))\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and we can plot it with with {func}`sympy.plot <sympy.plotting.plot.plot>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sy.plot(\n",
    "    plotted_intensity,\n",
    "    (m, 2 * pi0.mass, 2.5),\n",
    "    axis_center=(2 * pi0.mass, 0),\n",
    "    xlabel=fR\"$m(\\pi^{0}\\pi^{0})$\",\n",
    "    ylabel=\"$I$\",\n",
    "    backend=\"matplotlib\",\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expression itself looks like this (after some rounding of the {class}`float` values in this expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from sympy import preorder_traversal\n",
    "\n",
    "rounded_intensity = plotted_intensity\n",
    "rounded_intensity = rounded_intensity.subs({sy.sqrt(10): sy.sqrt(10).n()})\n",
    "for a in preorder_traversal(rounded_intensity):\n",
    "    if isinstance(a, sy.Float):\n",
    "        rounded_intensity = rounded_intensity.subs(a, round(a, 2))\n",
    "rounded_intensity"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}