{
 "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": {
    "tags": []
   },
   "source": [
    "# Dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{{ run_interactive }}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the dynamic terms in an amplitude model are set to $1$ by the {class}`.HelicityAmplitudeBuilder`. The method {meth}`.set_dynamics` can then be used to set dynamics lineshapes for specific resonances. The {mod}`.dynamics.builder` module provides some tools to set standard lineshapes (see below), but it is also possible to set {doc}`custom dynamics </usage/dynamics/custom>`.\n",
    "\n",
    "The standard lineshapes provided by AmpForm are illustrated below. For more info, have a look at the following pages:\n",
    "\n",
    "```{toctree}\n",
    ":maxdepth: 2\n",
    "dynamics/custom\n",
    "dynamics/analytic-continuation\n",
    "dynamics/k-matrix\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "import logging\n",
    "import warnings\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import mpl_interactions.ipyplot as iplt\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from IPython.display import Math\n",
    "\n",
    "import symplot\n",
    "\n",
    "logging.basicConfig()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Form factor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AmpForm uses Blatt-Weisskopf functions $B_L$ as _barrier factors_ (also called _form factors_, see {class}`.BlattWeisskopfSquared`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import BlattWeisskopfSquared\n",
    "\n",
    "L = sp.Symbol(\"L\", integer=True)\n",
    "z = sp.Symbol(\"z\", real=True)\n",
    "ff2 = BlattWeisskopfSquared(L, z)\n",
    "Math(sp.multiline_latex(ff2, ff2.doit(), environment=\"eqnarray\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Blatt-Weisskopf form factor is used to 'dampen' the breakup-momentum at threshold and when going to infinity. A usual choice for $z$ is therefore $z=q^2d^2$ with $q^2$ the {class}`.BreakupMomentumSquared` and $d$ the impact parameter (also called meson radius):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import BreakupMomentumSquared\n",
    "\n",
    "m, m_a, m_b, d = sp.symbols(\"m, m_a, m_b, d\")\n",
    "s = m**2\n",
    "q_squared = BreakupMomentumSquared(s, m_a, m_b)\n",
    "ff2 = BlattWeisskopfSquared(L, z=q_squared * d**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "np_blatt_weisskopf, sliders = symplot.prepare_sliders(\n",
    "    plot_symbol=m,\n",
    "    expression=ff2.doit(),\n",
    ")\n",
    "np_breakup_momentum = sp.lambdify((m, L, d, m_a, m_b), q_squared.doit())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "plot_domain = np.linspace(0.01, 4, 500)\n",
    "sliders.set_ranges(\n",
    "    L=(0, 8),\n",
    "    m_a=(0, 2, 200),\n",
    "    m_b=(0, 2, 200),\n",
    "    d=(0, 5),\n",
    ")\n",
    "sliders.set_values(\n",
    "    L=1,\n",
    "    m_a=0.3,\n",
    "    m_b=0.2,\n",
    "    d=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 5), tight_layout=True)\n",
    "ax.set_xlabel(\"$m$\")\n",
    "ax.axhline(0, c=\"black\", linewidth=0.5)\n",
    "controls = iplt.plot(\n",
    "    plot_domain,\n",
    "    np_blatt_weisskopf,\n",
    "    **sliders,\n",
    "    ylim=\"auto\",\n",
    "    label=R\"$B_L^2\\left(q(m)\\right)$\",\n",
    "    linestyle=\"dotted\",\n",
    "    ax=ax,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain,\n",
    "    np_breakup_momentum,\n",
    "    controls=controls,\n",
    "    ylim=\"auto\",\n",
    "    label=\"$q^2(m)$\",\n",
    "    linestyle=\"dashed\",\n",
    "    ax=ax,\n",
    ")\n",
    "iplt.title(\n",
    "    \"Effect of Blatt-Weisskopf factor\\n$L={L}, m_a={m_a:.1f}, m_b={m_b:.1f}$\",\n",
    "    controls=controls,\n",
    ")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "if STATIC_WEB_PAGE:\n",
    "    from IPython.display import SVG\n",
    "\n",
    "    output_file = \"blatt-weisskopf.svg\"\n",
    "    plt.savefig(output_file)\n",
    "    display(SVG(output_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Relativistic Breit-Wigner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AmpForm has two types of relativistic Breit-Wigner functions. Both are compared below ― for more info, see the links to the API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### _Without_ form factor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 'normal' {func}`.relativistic_breit_wigner` looks as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import relativistic_breit_wigner\n",
    "\n",
    "m, m0, w0 = sp.symbols(\"m, m0, Gamma0\")\n",
    "rel_bw = relativistic_breit_wigner(s=m**2, mass0=m0, gamma0=w0)\n",
    "rel_bw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _With_ form factor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The relativistic Breit-Wigner can be adapted slightly, so that its amplitude goes to zero at threshold ($m_0 = m_a + m_b$) and that it becomes normalizable. This is done with {ref}`form factors <usage/dynamics:Form factor>` and can be obtained with the function {func}`.relativistic_breit_wigner_with_ff`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import PhaseSpaceFactor  # noqa: F401\n",
    "from ampform.dynamics import BreakupMomentumSquared, ComplexSqrt\n",
    "\n",
    "\n",
    "def breakup_momentum(s: sp.Symbol, m_a: sp.Symbol, m_b: sp.Symbol) -> sp.Expr:\n",
    "    q_squared = BreakupMomentumSquared(s, m_a, m_b)\n",
    "    return ComplexSqrt(q_squared)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import (\n",
    "    PhaseSpaceFactorAnalytic,\n",
    "    relativistic_breit_wigner_with_ff,\n",
    ")\n",
    "\n",
    "rel_bw_with_ff = relativistic_breit_wigner_with_ff(\n",
    "    s=s,\n",
    "    mass0=m0,\n",
    "    gamma0=w0,\n",
    "    m_a=m_a,\n",
    "    m_b=m_b,\n",
    "    angular_momentum=L,\n",
    "    meson_radius=1,\n",
    "    phsp_factor=PhaseSpaceFactorAnalytic,\n",
    ")\n",
    "rel_bw_with_ff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, $\\Gamma(m)$ is the {class}`.EnergyDependentWidth` (also called running width or mass-dependent width), defined as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import EnergyDependentWidth\n",
    "\n",
    "L = sp.Symbol(\"L\", integer=True)\n",
    "width = EnergyDependentWidth(\n",
    "    s=s,\n",
    "    mass0=m0,\n",
    "    gamma0=w0,\n",
    "    m_a=m_a,\n",
    "    m_b=m_b,\n",
    "    angular_momentum=L,\n",
    "    meson_radius=1,\n",
    "    phsp_factor=PhaseSpaceFactorAnalytic,\n",
    ")\n",
    "Math(sp.multiline_latex(width, width.evaluate(), environment=\"eqnarray\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is possible to choose different formulations for the phase space factor $\\rho$, see {doc}`/usage/dynamics/analytic-continuation`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Analytic continuation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following shows the effect of {doc}`/usage/dynamics/analytic-continuation` a on relativistic Breit-Wigner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from ampform.dynamics import PhaseSpaceFactorComplex\n",
    "\n",
    "# Two types of relativistic Breit-Wigners\n",
    "rel_bw_with_ff = relativistic_breit_wigner_with_ff(\n",
    "    s=s,\n",
    "    mass0=m0,\n",
    "    gamma0=w0,\n",
    "    m_a=m_a,\n",
    "    m_b=m_b,\n",
    "    angular_momentum=L,\n",
    "    meson_radius=d,\n",
    "    phsp_factor=PhaseSpaceFactorComplex,\n",
    ")\n",
    "rel_bw_with_ff_ac = relativistic_breit_wigner_with_ff(\n",
    "    s=s,\n",
    "    mass0=m0,\n",
    "    gamma0=w0,\n",
    "    m_a=m_a,\n",
    "    m_b=m_b,\n",
    "    angular_momentum=L,\n",
    "    meson_radius=d,\n",
    "    phsp_factor=PhaseSpaceFactorAnalytic,\n",
    ")\n",
    "\n",
    "# Lambdify\n",
    "np_rel_bw_with_ff, sliders = symplot.prepare_sliders(\n",
    "    plot_symbol=m,\n",
    "    expression=rel_bw_with_ff.doit(),\n",
    ")\n",
    "np_rel_bw_with_ff_ac = sp.lambdify(\n",
    "    args=(m, w0, L, d, m0, m_a, m_b),\n",
    "    expr=rel_bw_with_ff_ac.doit(),\n",
    ")\n",
    "np_rel_bw = sp.lambdify(\n",
    "    args=(m, w0, L, d, m0, m_a, m_b),\n",
    "    expr=rel_bw.doit(),\n",
    ")\n",
    "\n",
    "# Set sliders\n",
    "plot_domain = np.linspace(start=0, stop=4, num=500)\n",
    "sliders.set_ranges(\n",
    "    m0=(0, 4, 200),\n",
    "    Gamma0=(0, 1, 100),\n",
    "    L=(0, 8),\n",
    "    m_a=(0, 2, 200),\n",
    "    m_b=(0, 2, 200),\n",
    "    d=(0, 5),\n",
    ")\n",
    "sliders.set_values(\n",
    "    m0=1.8,\n",
    "    Gamma0=0.6,\n",
    "    L=1,\n",
    "    m_a=0.6,\n",
    "    m_b=0.6,\n",
    "    d=1,\n",
    ")\n",
    "\n",
    "fig, axes = plt.subplots(\n",
    "    nrows=2,\n",
    "    figsize=(8, 6),\n",
    "    sharex=True,\n",
    "    #     gridspec_kw={\"height_ratios\": [1, 1.8]},\n",
    ")\n",
    "ax_ff, ax_ac = axes\n",
    "ax_ac.set_xlabel(\"$m$\")\n",
    "for ax in axes:\n",
    "    ax.axhline(0, c=\"gray\", linewidth=0.5)\n",
    "\n",
    "rho_c = PhaseSpaceFactorComplex(m, m_a, m_b)\n",
    "ax_ff.set_title(\n",
    "    f\"'Complex' phase space factor: ${sp.latex(rho_c.evaluate())}$\"\n",
    ")\n",
    "ax_ac.set_title(\"Analytic continuation of the phase space factor\")\n",
    "\n",
    "ylim = \"auto\"  # (-0.6, 1.2)\n",
    "controls = iplt.plot(\n",
    "    plot_domain,\n",
    "    lambda *args, **kwargs: np_rel_bw_with_ff(*args, **kwargs).real,\n",
    "    label=\"real\",\n",
    "    **sliders,\n",
    "    ylim=ylim,\n",
    "    ax=ax_ff,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain.real,\n",
    "    lambda *args, **kwargs: np_rel_bw_with_ff(*args, **kwargs).imag,\n",
    "    label=\"imaginary\",\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    ax=ax_ff,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain.real,\n",
    "    lambda *args, **kwargs: np.abs(np_rel_bw_with_ff(*args, **kwargs)) ** 2,\n",
    "    label=\"absolute\",\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    ax=ax_ff,\n",
    "    c=\"black\",\n",
    "    linestyle=\"dotted\",\n",
    ")\n",
    "\n",
    "\n",
    "iplt.plot(\n",
    "    plot_domain.real,\n",
    "    lambda *args, **kwargs: np_rel_bw_with_ff_ac(*args, **kwargs).real,\n",
    "    label=\"real\",\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    ax=ax_ac,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain.real,\n",
    "    lambda *args, **kwargs: np_rel_bw_with_ff_ac(*args, **kwargs).imag,\n",
    "    label=\"imaginary\",\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    ax=ax_ac,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain.real,\n",
    "    lambda *args, **kwargs: np.abs(np_rel_bw_with_ff_ac(*args, **kwargs)) ** 2,\n",
    "    label=\"absolute\",\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    ax=ax_ac,\n",
    "    c=\"black\",\n",
    "    linestyle=\"dotted\",\n",
    ")\n",
    "\n",
    "for ax in axes:\n",
    "    iplt.axvline(\n",
    "        controls[\"m0\"],\n",
    "        ax=ax,\n",
    "        c=\"red\",\n",
    "        label=f\"${sp.latex(m0)}$\",\n",
    "        alpha=0.3,\n",
    "    )\n",
    "    iplt.axvline(\n",
    "        lambda m_a, m_b, **kwargs: m_a + m_b,\n",
    "        controls=controls,\n",
    "        ax=ax,\n",
    "        c=\"black\",\n",
    "        alpha=0.3,\n",
    "        label=f\"${sp.latex(m_a)} + {sp.latex(m_b)}$\",\n",
    "    )\n",
    "ax_ac.legend(loc=\"upper right\")\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "if STATIC_WEB_PAGE:\n",
    "    from IPython.display import SVG\n",
    "\n",
    "    output_file = \"relativistic-breit-wigner-with-form-factor.svg\"\n",
    "    plt.savefig(output_file)\n",
    "    display(SVG(output_file))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}