{
 "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": [
    "# Custom dynamics"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "from typing import Dict, Tuple\n",
    "\n",
    "import graphviz\n",
    "import sympy as sp\n",
    "\n",
    "import expertsystem as es\n",
    "from expertsystem.amplitude.dynamics.builder import (\n",
    "    ResonanceDynamicsBuilder,\n",
    "    TwoBodyKinematicVariableSet,\n",
    "    create_relativistic_breit_wigner,\n",
    ")\n",
    "from expertsystem.reaction.particle import Particle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by generating allowed transitions for a simple decay channel, just like in {doc}`/usage/amplitude`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = es.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [+1]),\n",
    "    final_state=[(\"gamma\", [+1]), \"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": [
    "dot = es.io.asdot(result, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, create a {class}`.HelicityAmplitudeBuilder` using {func}`.get_builder`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_builder = es.amplitude.get_builder(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In {doc}`/usage/amplitude`, we used {meth}`.set_dynamics` with some standard lineshape builders from the {mod}`.builder` module. These builders have a signature that follows the {class}`.ResonanceDynamicsBuilder` {class}`~typing.Protocol`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(inspect.getsource(ResonanceDynamicsBuilder))\n",
    "print(inspect.getsource(create_relativistic_breit_wigner))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A function that behaves like a {class}`.ResonanceDynamicsBuilder` should return a {class}`tuple` of some {class}`~sympy.core.expr.Expr` (which formulates your lineshape) and a {class}`dict` of {class}`~sympy.core.symbol.Symbol`s to some suggested initial values. This signature is required so that {meth}`.set_dynamics` knows how to extract the correct symbol names and their suggested initial values from a {class}`.StateTransitionGraph`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {class}`~sympy.core.expr.Expr` you use for the lineshape can be anything. Here, we use a Gaussian function and wrap it in a function. As you can see, this function stands on its own, independent of the {mod}`expertsystem`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_dynamics(x: sp.Symbol, mu: sp.Symbol, sigma: sp.Symbol) -> sp.Expr:\n",
    "    return sp.exp(-((x - mu) ** 2) / sigma ** 2 / 2) / (\n",
    "        sigma * sp.sqrt(2 * sp.pi)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, mu, sigma = sp.symbols(\"x mu sigma\")\n",
    "sp.plot(my_dynamics(x, 0, 1), (x, -3, 3), axis_center=(0, 0))\n",
    "my_dynamics(x, mu, sigma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now follow the example of the {func}`.create_relativistic_breit_wigner` to create a builder for this custom lineshape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_my_dynamics(\n",
    "    particle: Particle, variable_pool: TwoBodyKinematicVariableSet\n",
    ") -> Tuple[sp.Expr, Dict[sp.Symbol, float]]:\n",
    "    res_mass = sp.Symbol(f\"m_{particle.name}\")\n",
    "    res_width = sp.Symbol(f\"sigma_{particle.name}\")\n",
    "    return (\n",
    "        my_dynamics(\n",
    "            x=variable_pool.in_edge_inv_mass,\n",
    "            mu=res_mass,\n",
    "            sigma=res_width,\n",
    "        ),\n",
    "        {res_mass: particle.mass, res_width: particle.width},\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, just like in {ref}`usage/amplitude:Build SymPy expression`, it's simply a matter of plugging this builder into {meth}`.set_dynamics` and we can {meth}`~.HelicityAmplitudeBuilder.generate` a model with this custom lineshape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in result.get_intermediate_particles().names:\n",
    "    model_builder.set_dynamics(name, create_my_dynamics)\n",
    "model = model_builder.generate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, the {attr}`.HelicityModel.parameter_defaults` section has been updated with the some additional parameters for the custom parameter and there corresponding suggested initial values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.parameter_defaults"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's quickly have a look what this lineshape looks like. First, check which {class}`~sympy.core.symbol.Symbol`s remain once we replace the parameters with their suggested initial values. These are the kinematic variables of the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr = model.expression.doit().subs(model.parameter_defaults)\n",
    "free_symbols = tuple(sorted(expr.free_symbols, key=lambda s: s.name))\n",
    "free_symbols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To create an invariant mass distribution, we should integrate out the $\\theta$ angle. This can be done with {func}`~sympy.integrals.integrals.integrate`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m, theta = free_symbols\n",
    "integrated_expr = sp.integrate(\n",
    "    expr,\n",
    "    (theta, 0, sp.pi),\n",
    "    meijerg=True,\n",
    "    conds=\"piecewise\",\n",
    "    risch=None,\n",
    "    heurisch=None,\n",
    "    manual=None,\n",
    ")\n",
    "integrated_expr.n(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, here is the resulting expression as a function of the invariant mass, with **custom dynamics**!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, x2 = 0.6, 1.9\n",
    "sp.plot(integrated_expr, (m, x1, x2), axis_center=(x1, 0));"
   ]
  }
 ],
 "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
}