docs/report/003.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "physics",
     "sympy"
    ]
   },
   "source": [
    "::::{margin}\n",
    ":::{card} Chew-Mandelstam dispersion integrals\n",
    "TR-003\n",
    "^^^\n",
    "This report formulates a symbolic dispersion integral to approach the left-hand cut in the form factor for arbitrary angular momentum. The integral is evaluated with SciPy's [`quad_vec`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad_vec.html).\n",
    "+++\n",
    "🚧&nbsp;[ampform#265](https://github.com/ComPWA/ampform/issues/265)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chew-Mandelstam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q ampform==0.14.10 black==23.12.1 mpl-interactions==0.24.1 scipy==1.12.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "This report is an attempt formulate the Chew-Mandelstam function described in [PDG2023, §50.3.3 (Resonances), pp.14–15](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15) with [SymPy](https://docs.sympy.org), so that it can be implemented in [AmpForm](https://ampform.rtfd.io)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Import Python libraries"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "import inspect\n",
    "import os\n",
    "import warnings\n",
    "from functools import partial\n",
    "from typing import Any\n",
    "\n",
    "import black\n",
    "import matplotlib.pyplot as plt\n",
    "import mpl_interactions.ipyplot as iplt\n",
    "import numpy as np\n",
    "import qrules\n",
    "import sympy as sp\n",
    "from ampform.dynamics import (\n",
    "    BlattWeisskopfSquared,\n",
    "    BreakupMomentumSquared,\n",
    "    PhaseSpaceFactor,\n",
    "    PhaseSpaceFactorComplex,\n",
    ")\n",
    "from ampform.io import aslatex\n",
    "from ampform.sympy import unevaluated\n",
    "from ampform.sympy.math import ComplexSqrt\n",
    "from IPython.display import SVG, Markdown, Math, display\n",
    "from ipywidgets import FloatSlider\n",
    "from scipy.integrate import quad_vec\n",
    "from sympy.printing.pycode import _unpack_integral_limits\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "PDG = qrules.load_pdg()\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "if \"COLAB_RELEASE_TAG\" in os.environ:\n",
    "    import subprocess\n",
    "\n",
    "    from google.colab import output\n",
    "\n",
    "    output.enable_custom_widget_manager()\n",
    "    subprocess.run(\"pip install -q ipympl\".split(), check=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## S-wave"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "As can be seen in Eq. (50.44) on [PDG2023, §Resonances, p.15](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15), the Chew-Mandelstam function $\\Sigma_a$ for a particle $a$ decaying to particles $1, 2$ has a simple form for angular momentum $L=0$ ($S$-wave):\n",
    "\n",
    ":::{math}\n",
    ":class: full-width\n",
    ":label: chew-mandelstam\n",
    "\\Sigma_a(s) = \\frac{1}{16\\pi^2}\n",
    "\\left[\n",
    " \\frac{2q_a}{\\sqrt{s}}\n",
    " \\log\\frac{m_1^2+m_2^2-s+2\\sqrt{s}q_a}{2m_1m_2}\n",
    " - \\left(m_1^2-m_2^2\\right)\n",
    " \\left(\\frac{1}{s}-\\frac{1}{(m_1+m_2)^2}\\right)\n",
    " \\log\\frac{m_1}{m_2}\n",
    "\\right]\n",
    ":::\n",
    "\n",
    "The only question is how to deal with negative values for the squared break-up momentum $q_a^2$. Here, we will use AmpForm's {class}`~ampform.sympy.math.ComplexSqrt`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "z = sp.Symbol(\"z\")\n",
    "sqrt_expr = ComplexSqrt(z)\n",
    "Math(aslatex({sqrt_expr: sqrt_expr.get_definition()}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input",
     "scroll-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "@unevaluated\n",
    "class BreakupMomentum(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    _latex_repr_ = R\"q_a\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        q_squared = BreakupMomentumSquared(s, m1, m2)\n",
    "        return ComplexSqrt(q_squared)\n",
    "\n",
    "\n",
    "@unevaluated\n",
    "class ChewMandelstamSWave(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        # evaluate=False in order to keep same style as PDG\n",
    "        s, m1, m2 = self.args\n",
    "        q = BreakupMomentum(s, m1, m2)\n",
    "        left_term = sp.Mul(\n",
    "            2 * q / sp.sqrt(s),\n",
    "            sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)),\n",
    "            evaluate=False,\n",
    "        )\n",
    "        right_term = (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n",
    "        return sp.Mul(\n",
    "            1 / (16 * sp.pi**2),\n",
    "            left_term - right_term,\n",
    "            evaluate=False,\n",
    "        )\n",
    "\n",
    "\n",
    "s, m1, m2 = sp.symbols(\"s m1 m2 \")\n",
    "cm_expr = ChewMandelstamSWave(s, m1, m2)\n",
    "q_expr = BreakupMomentum(s, m1, m2)\n",
    "Math(aslatex({e: e.doit(deep=False) for e in [cm_expr, q_expr]}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "It should be noted that this equation is not well-defined along the real axis, that is, for $\\mathrm{Im}(s) = 0$. For this reason, we split $s$ into a real part $s'$ with a small imaginary offset (the PDG indicates this with $s+0i$). We parametrized this imaginary offset with $\\epsilon$, and for the interactive plot, we do so with a power of $10$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "s_prime = sp.Symbol(R\"s^{\\prime}\", real=True)\n",
    "epsilon = sp.Symbol(\"epsilon\", positive=True)\n",
    "s_plus = s_prime + sp.I * sp.Pow(10, -epsilon)\n",
    "Math(Rf\"{sp.latex(s)} \\to {sp.latex(s_plus)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "source": [
    "We are now ready to use express the symbolic expressions above as a numerical function. For comparison, we will plot the Chew-Mandelstam function for $S$-waves next to AmpForm's {class}`~ampform.dynamics.phasespace.PhaseSpaceFactorComplex`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "rho_expr = PhaseSpaceFactorComplex(s, m1, m2)\n",
    "Math(aslatex({rho_expr: rho_expr.doit(deep=False)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "symbols = (s_prime, m1, m2, epsilon)\n",
    "cm_func = sp.lambdify(symbols, cm_expr.doit().subs(s, s_plus))\n",
    "rho_func = sp.lambdify(symbols, rho_expr.doit().subs(s, s_plus))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "As starting values for the interactive plot, we assume $\\pi\\eta$ scattering (just like in the PDG section) and use their masses as values for $m_1$ and $m_1$, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Define sliders and plot range"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "s_min, s_max = -0.15, 1.4\n",
    "plot_domain = np.linspace(s_min, s_max, num=500)\n",
    "\n",
    "m1_val = PDG[\"pi0\"].mass\n",
    "m2_val = PDG[\"eta\"].mass\n",
    "sliders = dict(\n",
    "    m1=FloatSlider(\n",
    "        description=str(m1),\n",
    "        min=0,\n",
    "        max=1.2,\n",
    "        value=m1_val,\n",
    "    ),\n",
    "    m2=FloatSlider(\n",
    "        description=str(m2),\n",
    "        min=0,\n",
    "        max=1.2,\n",
    "        value=m2_val,\n",
    "    ),\n",
    "    epsilon=FloatSlider(\n",
    "        description=str(epsilon),\n",
    "        min=0,\n",
    "        max=8,\n",
    "        value=4,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-output",
     "scroll-input",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(ncols=2, figsize=(11, 4.5), tight_layout=True)\n",
    "ax1, ax2 = axes\n",
    "for ax in axes:\n",
    "    ax.axhline(0, linewidth=0.5, c=\"black\")\n",
    "\n",
    "real_style = {\"label\": \"Real part\", \"c\": \"black\", \"linestyle\": \"dashed\"}\n",
    "imag_style = {\"label\": \"Imag part\", \"c\": \"red\"}\n",
    "threshold_style = {\"label\": R\"$s_\\mathrm{thr}$\", \"c\": \"grey\", \"linewidth\": 0.5}\n",
    "\n",
    "xlim = (plot_domain.min(), plot_domain.max())\n",
    "ylim = (-1, +1)\n",
    "y_factor = 16 * np.pi\n",
    "controls = iplt.axvline(\n",
    "    lambda *args, **kwargs: (kwargs[\"m1\"] + kwargs[\"m2\"]) ** 2,\n",
    "    **sliders,\n",
    "    ax=ax1,\n",
    "    **threshold_style,\n",
    ")\n",
    "iplt.axvline(\n",
    "    lambda *args, **kwargs: (kwargs[\"m1\"] + kwargs[\"m2\"]) ** 2,\n",
    "    controls=controls,\n",
    "    ax=ax2,\n",
    "    **threshold_style,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain,\n",
    "    lambda *args, **kwargs: (y_factor * 1j * rho_func(*args, **kwargs)).real,\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    xlim=xlim,\n",
    "    alpha=0.7,\n",
    "    ax=ax1,\n",
    "    **real_style,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain,\n",
    "    lambda *args, **kwargs: (y_factor * 1j * rho_func(*args, **kwargs)).imag,\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    xlim=xlim,\n",
    "    alpha=0.7,\n",
    "    ax=ax1,\n",
    "    **imag_style,\n",
    ")\n",
    "\n",
    "iplt.plot(\n",
    "    plot_domain,\n",
    "    lambda *args, **kwargs: y_factor * cm_func(*args, **kwargs).real,\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    xlim=xlim,\n",
    "    alpha=0.7,\n",
    "    ax=ax2,\n",
    "    **real_style,\n",
    ")\n",
    "iplt.plot(\n",
    "    plot_domain,\n",
    "    lambda *args, **kwargs: y_factor * cm_func(*args, **kwargs).imag,\n",
    "    controls=controls,\n",
    "    ylim=ylim,\n",
    "    xlim=xlim,\n",
    "    alpha=0.7,\n",
    "    ax=ax2,\n",
    "    **imag_style,\n",
    ")\n",
    "\n",
    "for ax in axes:\n",
    "    ax.legend(loc=\"lower right\")\n",
    "    ax.set_xticks(np.arange(0, 1.21, 0.3))\n",
    "    ax.set_yticks(np.arange(-1, 1.1, 0.5))\n",
    "    ax.set_xlim()\n",
    "    ax.set_xlabel(\"$s$ (GeV$^2$)\")\n",
    "\n",
    "ax1.set_ylabel(R\"$16\\pi \\; i\\rho(s)$\")\n",
    "ax2.set_ylabel(R\"$16\\pi \\; \\Sigma(s)$\")\n",
    "ax1.set_title(R\"Complex phase space factor $\\rho$\")\n",
    "ax2.set_title(\"Chew-Mandelstam $S$-wave ($L=0$)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "if STATIC_WEB_PAGE:\n",
    "    filename = \"003/chew-mandelstam-s-wave.svg\"\n",
    "    os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
    "    fig.savefig(filename)\n",
    "    display(*sliders.values(), SVG(filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ":::{tip}\n",
    "Compare the plots above with Figure 50.6 on [PDG2023, §Resonances, p.16](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15).\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## General dispersion integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "For higher angular momenta, the PDG notes that one has to compute the dispersion integral given by Eq. (50.44) on [PDG2023, §Resonances, p.15](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15):\n",
    "\n",
    "$$\n",
    "\\Sigma_a(s+0i) =\n",
    "    \\frac{s-s_{\\mathrm{thr}_a}}{\\pi}\n",
    "    \\int^\\infty_{s_{\\mathrm{thr}_a}} \\frac{\n",
    "        \\rho_a(s')n_a^2(s')\n",
    "    }{\n",
    "        (s' - s_{\\mathrm{thr}_a})(s'-s-i0)\n",
    "    }\n",
    "    \\mathop{}\\!\\mathrm{d}s'\n",
    "$$ (dispersion-integral)\n",
    "\n",
    "Equation {eq}`chew-mandelstam` is the analytic solution for $L=0$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "From Equations (50.33–34) on [PDG2023, §Resonances, p.12](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=12), it can be deduced that the function $n_a^2$ is the same as AmpForm's {class}`~ampform.dynamics.BlattWeisskopfSquared` (note that this function is normalized, whereas the PDG's $F_j$ function has $1$ in the nominator). For this reason, we simply use {class}`~ampform.dynamics.BlattWeisskopfSquared` for the definition of $n_a^2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "@unevaluated\n",
    "class FormFactor(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    L: Any\n",
    "    q0: Any = 1\n",
    "    _latex_repr_ = R\"n_a^2\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2, L, q0 = self.args\n",
    "        q_squared = BreakupMomentumSquared(s, m1, m2)\n",
    "        return BlattWeisskopfSquared(\n",
    "            z=q_squared / (q0**2),\n",
    "            angular_momentum=L,\n",
    "        )\n",
    "\n",
    "\n",
    "L = sp.Symbol(\"L\", integer=True, positive=True)\n",
    "q0 = sp.Symbol(\"q0\", real=True)\n",
    "na_expr = FormFactor(s, m1, m2, L, q0)\n",
    "bl_expr = BlattWeisskopfSquared(z, L)\n",
    "Math(aslatex({e: e.doit(deep=False) for e in [na_expr, bl_expr]}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "For $\\rho_a$, we use AmpForm's {class}`~ampform.dynamics.phasespace.PhaseSpaceFactor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Math(aslatex({rho_expr: rho_expr.doit(deep=False)}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The symbolic integrand is then formulated as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    }
   },
   "outputs": [],
   "source": [
    "s_thr = (m1 + m2) ** 2\n",
    "integrand = (PhaseSpaceFactor(s_prime, m1, m2) * FormFactor(s_prime, m1, m2, L, q0)) / (\n",
    "    (s_prime - s_thr) * (s_prime - s - epsilon * sp.I)\n",
    ")\n",
    "integrand"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Next, we {func}`~sympy.utilities.lambdify.lambdify` this integrand to a {mod}`numpy` expression so that we can integrate it efficiently:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "integrand_func = sp.lambdify(\n",
    "    args=(s_prime, s, L, epsilon, m1, m2, q0),\n",
    "    expr=integrand.doit(),\n",
    "    modules=\"numpy\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ":::{note}\n",
    "Integrals can be expressed symbolically with SymPy, with some caveats. See {ref}`report/016:SymPy integral`.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "As discussed in [TR-016](016.ipynb), {func}`scipy.integrate.quad` cannot integrate over complex-valued functions, but {func}`scipy.integrate.quad_vec` can. For comparison, we now compute this integral for a few values of $L>0$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "s_domain = np.linspace(s_min, s_max, num=50)\n",
    "max_L = 3\n",
    "l_values = list(range(1, max_L + 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "It is handy to store the numerical results of each dispersion integral in a {obj}`dict` with $L$ as keys:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    }
   },
   "outputs": [],
   "source": [
    "s_thr_val = float(s_thr.subs({m1: m1_val, m2: m2_val}))\n",
    "integral_values = {\n",
    "    l_val: quad_vec(\n",
    "        lambda x: integrand_func(\n",
    "            x,\n",
    "            s=s_domain,\n",
    "            L=l_val,\n",
    "            epsilon=1e-3,\n",
    "            m1=m1_val,\n",
    "            m2=m2_val,\n",
    "            q0=1.0,\n",
    "        ),\n",
    "        a=s_thr_val,\n",
    "        b=np.inf,\n",
    "    )[0]\n",
    "    for l_val in tqdm(l_values, desc=\"Evaluating integrals\")\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Finally, as can be seen from Eq. {eq}`dispersion-integral`, the resulting values from the integral have to be shifted with a factor $\\frac{s-s_{\\mathrm{thr}_a}}{\\pi}$ to get $\\Sigma_a$. We also scale the values with $16\\pi$ so that it can be compared with the plot generated in {ref}`report/003:S-wave`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sigma = {\n",
    "    l_val: (s_domain - s_thr_val) / np.pi * integral_values[l_val] for l_val in l_values\n",
    "}\n",
    "sigma_scaled = {l_val: 16 * np.pi * sigma[l_val] for l_val in l_values}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-output",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(\n",
    "    nrows=len(l_values),\n",
    "    sharex=True,\n",
    "    figsize=(5, 2.5 * len(l_values)),\n",
    "    tight_layout=True,\n",
    ")\n",
    "fig.suptitle(f\"Dispersion integrals for $m_1={m1_val:.2f}, m_2={m2_val:.2f}$\")\n",
    "for ax, l_val in zip(axes, l_values):\n",
    "    ax.axhline(0, linewidth=0.5, c=\"black\")\n",
    "    ax.axvline(s_thr_val, **threshold_style)\n",
    "    ax.plot(s_domain, sigma_scaled[l_val].real, **real_style)\n",
    "    ax.plot(s_domain, sigma_scaled[l_val].imag, **imag_style)\n",
    "    ax.set_title(f\"$L = {l_val}$\")\n",
    "    ax.set_ylabel(R\"$16\\pi \\; \\Sigma(s)$\")\n",
    "axes[-1].set_xlabel(\"$s$ (GeV$^2$)\")\n",
    "axes[0].legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "if STATIC_WEB_PAGE:\n",
    "    filename = \"003/chew-mandelstam-l-non-zero.svg\"\n",
    "    os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
    "    fig.savefig(filename)\n",
    "    display(SVG(filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## SymPy expressions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following, we attempt to implement Equation {eq}`dispersion-integral` using {ref}`report/016:SymPy integral`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "mystnb": {
     "code_prompt_show": "Symbolic integral definition"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "class UnevaluatableIntegral(sp.Integral):\n",
    "    abs_tolerance = 1e-5\n",
    "    rel_tolerance = 1e-5\n",
    "    limit = 50\n",
    "    dummify = True\n",
    "\n",
    "    def doit(self, **hints):\n",
    "        args = [arg.doit(**hints) for arg in self.args]\n",
    "        return self.func(*args)\n",
    "\n",
    "    def _numpycode(self, printer, *args):\n",
    "        integration_vars, limits = _unpack_integral_limits(self)\n",
    "        if len(limits) != 1 or len(integration_vars) != 1:\n",
    "            msg = f\"Cannot handle {len(limits)}-dimensional integrals\"\n",
    "            raise ValueError(msg)\n",
    "        x = integration_vars[0]\n",
    "        a, b = limits[0]\n",
    "        expr = self.args[0]\n",
    "        if self.dummify:\n",
    "            dummy = sp.Dummy()\n",
    "            expr = expr.xreplace({x: dummy})\n",
    "            x = dummy\n",
    "        integrate_func = \"quad_vec\"\n",
    "        printer.module_imports[\"scipy.integrate\"].add(integrate_func)\n",
    "        return (\n",
    "            f\"{integrate_func}(lambda {printer._print(x)}: {printer._print(expr)},\"\n",
    "            f\" {printer._print(a)}, {printer._print(b)},\"\n",
    "            f\" epsabs={self.abs_tolerance}, epsrel={self.abs_tolerance},\"\n",
    "            f\" limit={self.limit})[0]\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Definition of the symbolic dispersion integral"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def dispersion_integral(\n",
    "    s,\n",
    "    m1,\n",
    "    m2,\n",
    "    angular_momentum,\n",
    "    meson_radius=1,\n",
    "    s_prime=sp.Symbol(\"x\", real=True),\n",
    "    epsilon=sp.Symbol(\"epsilon\", positive=True),\n",
    "):\n",
    "    s_thr = (m1 + m2) ** 2\n",
    "    q_squared = BreakupMomentumSquared(s_prime, m1, m2)\n",
    "    ff_squared = BlattWeisskopfSquared(\n",
    "        angular_momentum=L, z=q_squared * meson_radius**2\n",
    "    )\n",
    "    phsp_factor = PhaseSpaceFactor(s_prime, m1, m2)\n",
    "    return sp.Mul(\n",
    "        (s - s_thr) / sp.pi,\n",
    "        UnevaluatableIntegral(\n",
    "            (phsp_factor * ff_squared)\n",
    "            / (s_prime - s_thr)\n",
    "            / (s_prime - s - sp.I * epsilon),\n",
    "            (s_prime, s_thr, sp.oo),\n",
    "        ),\n",
    "        evaluate=False,\n",
    "    )\n",
    "\n",
    "\n",
    "integral_expr = dispersion_integral(s, m1, m2, angular_momentum=L, s_prime=s_prime)\n",
    "integral_expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ":::{warning}\n",
    "\n",
    "We have to keep track of the integration variable ($s'$ in Equation {eq}`dispersion-integral`), so that we don't run into trouble if we use {func}`~sympy.utilities.lambdify.lambdify` with common sub-expressions. The problem is that the integration variable _should not_ be extracted as a common sub-expression, otherwise the lambdified {func}`scipy.integrate.quad_vec` expression cannot handle vectorized input.\n",
    "\n",
    ":::\n",
    "\n",
    "To keep the function under the integral simple, we substitute angular momentum $L$ with a definite value before we lambdify:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "UnevaluatableIntegral.abs_tolerance = 1e-4\n",
    "UnevaluatableIntegral.rel_tolerance = 1e-4\n",
    "integral_s_wave_func = sp.lambdify(\n",
    "    [s, m1, m2, epsilon],\n",
    "    integral_expr.subs(L, 0).doit(),\n",
    "    # integration symbol should not be extracted as common sub-expression!\n",
    "    cse=partial(sp.cse, ignore=[s_prime], list=False),\n",
    ")\n",
    "integral_s_wave_func = np.vectorize(integral_s_wave_func)\n",
    "\n",
    "integral_p_wave_func = sp.lambdify(\n",
    "    [s, m1, m2, epsilon],\n",
    "    integral_expr.subs(L, 1).doit(),\n",
    "    cse=partial(sp.cse, ignore=[s_prime], list=False),\n",
    ")\n",
    "integral_p_wave_func = np.vectorize(integral_p_wave_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "src = inspect.getsource(integral_s_wave_func.pyfunc)\n",
    "src = f\"\"\"```python\n",
    "{black.format_str(src, mode=black.FileMode()).strip()}\n",
    "```\"\"\"\n",
    "Markdown(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "s_values = np.linspace(-0.15, 1.4, num=200)\n",
    "%time s_wave_values = integral_s_wave_func(s_values, m1_val, m2_val, epsilon=1e-5)\n",
    "%time p_wave_values = integral_p_wave_func(s_values, m1_val, m2_val, epsilon=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Note that the dispersion integral for $L=0$ indeed reproduces the same shape as in {ref}`report/003:S-wave`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "s_wave_values *= 16 * np.pi\n",
    "p_wave_values *= 16 * np.pi\n",
    "\n",
    "s_values = np.linspace(-0.15, 1.4, num=200)\n",
    "fig, axes = plt.subplots(nrows=2, figsize=(6, 7), sharex=True)\n",
    "ax1, ax2 = axes\n",
    "fig.suptitle(f\"Symbolic dispersion integrals for $m_1={m1_val:.2f}, m_2={m2_val:.2f}$\")\n",
    "for ax in axes:\n",
    "    ax.axhline(0, linewidth=0.5, c=\"black\")\n",
    "    ax.axvline(s_thr_val, **threshold_style)\n",
    "    ax.set_title(f\"$L = {l_val}$\")\n",
    "    ax.set_ylabel(R\"$16\\pi \\; \\Sigma(s)$\")\n",
    "axes[-1].set_xlabel(\"$s$ (GeV$^2$)\")\n",
    "\n",
    "ax1.set_title(\"$S$-wave ($L=0$)\")\n",
    "ax1.plot(s_values, s_wave_values.real, **real_style)\n",
    "ax1.plot(s_values, s_wave_values.imag, **imag_style)\n",
    "\n",
    "ax2.set_title(\"$P$-wave ($L=1$)\")\n",
    "ax2.plot(s_values, p_wave_values.real, **real_style)\n",
    "ax2.plot(s_values, p_wave_values.imag, **imag_style)\n",
    "\n",
    "ax1.legend()\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "if STATIC_WEB_PAGE:\n",
    "    filename = \"003/symbolic-chew-mandelstam.svg\"\n",
    "    os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
    "    fig.savefig(filename)\n",
    "    display(SVG(filename))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}