{
 "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. ampform==0.1.2\n",
    "%pip install -q ampform[doc,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": [
    "# Custom dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import graphviz\n",
    "import qrules\n",
    "import sympy as sp"
   ]
  },
  {
   "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": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "reaction = qrules.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=\"helicity\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction, 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": [
    "from ampform import get_builder\n",
    "\n",
    "model_builder = get_builder(reaction)"
   ]
  },
  {
   "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": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "\n",
    "from ampform.dynamics.builder import (\n",
    "    ResonanceDynamicsBuilder,\n",
    "    create_relativistic_breit_wigner,\n",
    ")\n",
    "\n",
    "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}`~qrules.transition.StateTransition`."
   ]
  },
  {
   "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 {mod}`ampform`:"
   ]
  },
  {
   "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": [
    "from typing import Dict, Tuple\n",
    "\n",
    "from qrules.particle import Particle\n",
    "\n",
    "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n",
    "\n",
    "\n",
    "def create_my_dynamics(\n",
    "    resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n",
    ") -> Tuple[sp.Expr, Dict[sp.Symbol, float]]:\n",
    "    res_mass = sp.Symbol(f\"m_{resonance.name}\")\n",
    "    res_width = sp.Symbol(f\"sigma_{resonance.name}\")\n",
    "    expression = my_dynamics(\n",
    "        x=variable_pool.incoming_state_mass,\n",
    "        mu=res_mass,\n",
    "        sigma=res_width,\n",
    "    )\n",
    "    parameter_defaults = {\n",
    "        res_mass: resonance.mass,\n",
    "        res_width: resonance.width,\n",
    "    }\n",
    "    return expression, parameter_defaults"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, just like in {ref}`usage/amplitude:Set dynamics`, it's simply a matter of plugging this builder into {meth}`.set_dynamics` and we can {meth}`~.HelicityAmplitudeBuilder.formulate` a model with this custom lineshape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in reaction.get_intermediate_particles().names:\n",
    "    model_builder.set_dynamics(name, create_my_dynamics)\n",
    "model = model_builder.formulate()"
   ]
  },
  {
   "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": [
    "display(*sorted(model.parameter_defaults, key=lambda p: p.name))"
   ]
  },
  {
   "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": {
    "tags": [
     "full-width"
    ]
   },
   "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 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}