docs/report/027.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "::::{margin}\n",
    ":::{card}  Visualization of the Riemann sheets for the two-channel $T$-matrix with one pole\n",
    "TR-027\n",
    "^^^\n",
    "Following **[TR-026](./026.ipynb)**, the Riemann sheets for the amplitude calculated within the $K$-matrix formalism for the two-channel case are visualized. The method of transitioning from the first physical sheet to the unphysical sheets is extended to the two dimensional case using [Eur. Phys. J. C (2023) 83:850](https://juser.fz-juelich.de/record/1017534/files/s10052-023-11953-6.pdf) in order to visualize the third and the fourth unphysical sheet.\n",
    "+++\n",
    "🚧&nbsp;[ampform#67](https://github.com/ComPWA/ampform/issues/67)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "# Coupled channel Riemann sheets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q ampform==0.15.0 plotly==5.18.0 sympy==1.12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Import Python libraries"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "from __future__ import annotations\n",
    "\n",
    "import os\n",
    "import warnings\n",
    "from typing import Any\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "import sympy as sp\n",
    "from ampform.io import aslatex\n",
    "from ampform.kinematics.phasespace import Kallen\n",
    "from ampform.sympy import unevaluated\n",
    "from IPython.display import Math, display\n",
    "from ipywidgets import widgets as w\n",
    "from plotly.colors import DEFAULT_PLOTLY_COLORS\n",
    "from plotly.subplots import make_subplots\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "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": "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,
    "tags": []
   },
   "source": [
    "## Expression definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "@unevaluated(real=False)\n",
    "class PhaseSpaceFactor(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    _latex_repr_ = R\"\\rho_{{{m1}, {m2}}}\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        return sp.sqrt((s - ((m1 + m2) ** 2)) * (s - (m1 - m2) ** 2) / s**2)\n",
    "\n",
    "\n",
    "s, m1, m2 = sp.symbols(\"s m1 m2\")\n",
    "rho_expr = PhaseSpaceFactor(s, m1, m2)\n",
    "Math(aslatex({rho_expr: rho_expr.doit(deep=False)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "@unevaluated(real=False)\n",
    "class PhaseSpaceFactorKallen(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    _latex_repr_ = R\"\\rho_{{{m1}, {m2}}}\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        return 2 * BreakupMomentum(s, m1, m2) / sp.sqrt(s)\n",
    "\n",
    "\n",
    "@unevaluated(real=False)\n",
    "class PhaseSpaceCM(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n",
    "\n",
    "\n",
    "@unevaluated(real=False)\n",
    "class ChewMandelstam(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",
    "        s, m1, m2 = self.args\n",
    "        q = BreakupMomentum(s, m1, m2)\n",
    "        return (\n",
    "            1\n",
    "            / (16 * sp.pi**2)\n",
    "            * (\n",
    "                (2 * q / sp.sqrt(s))\n",
    "                * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n",
    "                - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n",
    "            )\n",
    "        )\n",
    "\n",
    "\n",
    "@unevaluated(real=False)\n",
    "class BreakupMomentum(sp.Expr):\n",
    "    s: Any\n",
    "    m1: Any\n",
    "    m2: Any\n",
    "    _latex_repr_ = R\"q\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))\n",
    "\n",
    "\n",
    "s, m1, m2 = sp.symbols(\"s m1 m2\")\n",
    "rho_expr_kallen = PhaseSpaceFactorKallen(s, m1, m2)\n",
    "rho_cm_expr = PhaseSpaceCM(s, m1, m2)\n",
    "cm_expr = ChewMandelstam(s, m1, m2)\n",
    "q_expr = BreakupMomentum(s, m1, m2)\n",
    "kallen = Kallen(*sp.symbols(\"x:z\"))\n",
    "Math(\n",
    "    aslatex({\n",
    "        e: e.doit(deep=False)\n",
    "        for e in [rho_expr_kallen, rho_cm_expr, cm_expr, q_expr, kallen]\n",
    "    })\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "## Riemann sheet I"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "### Matrix definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "class DiagonalMatrix(sp.DiagonalMatrix):\n",
    "    def _latex(self, printer, *args):\n",
    "        return printer._print(self.args[0])\n",
    "\n",
    "\n",
    "n = 2\n",
    "I = sp.Identity(n)\n",
    "K = sp.MatrixSymbol(\"K\", n, n)\n",
    "CM = DiagonalMatrix(sp.MatrixSymbol(R\"\\rho^\\Sigma\", n, n))\n",
    "Math(aslatex({CM: CM.as_explicit()}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "T_I = (I - sp.I * K * CM).inv() * K\n",
    "T_I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "T_I_explicit = T_I.as_explicit()\n",
    "T_I_explicit[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "### Parametrization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Symbol definitions"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "s = sp.Symbol(\"s\")\n",
    "ma1 = sp.Symbol(\"m_{a1}\")\n",
    "mb1 = sp.Symbol(\"m_{b1}\")\n",
    "ma2 = sp.Symbol(\"m_{a2}\")\n",
    "mb2 = sp.Symbol(\"m_{b2}\")\n",
    "m0 = sp.Symbol(\"m0\")\n",
    "w0 = sp.Symbol(\"Gamma0\")\n",
    "g1 = sp.Symbol(R\"g^{0}_1\")\n",
    "g2 = sp.Symbol(R\"g^{0}_2\")\n",
    "symbols = sp.Tuple(s, ma1, mb1, ma2, mb2, m0, g1, g2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "k_expr_00 = (g1 * g1 * m0) / (s - m0**2)\n",
    "k_expr_10 = (g1 * g2 * m0) / (s - m0**2)\n",
    "k_expr_11 = (g2 * g2 * m0) / (s - m0**2)\n",
    "cm_expressions = {\n",
    "    K[0, 0]: k_expr_00,\n",
    "    K[1, 1]: k_expr_11,\n",
    "    K[0, 1]: k_expr_10,\n",
    "    K[1, 0]: k_expr_10,\n",
    "    CM[0, 0]: -PhaseSpaceCM(s, ma1, mb1),\n",
    "    CM[1, 1]: -PhaseSpaceCM(s, ma2, mb2),\n",
    "}\n",
    "Math(aslatex(cm_expressions))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "T_I_cm_expr = T_I_explicit.xreplace(cm_expressions)\n",
    "T_I_cm_expr[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "## Sheets II, III, and IV"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "In the case of two channels, there are four Riemann sheets. The first sheet ([Sheet I](#riemann-sheet-i)) is physical and three unphysical ones. The physical sheet is calculated using the analytic solution of the Chew-Mandelstam function.\n",
    "\n",
    "$$\n",
    "\\begin{eqnarray}\n",
    "\\operatorname{Disc}_{\\mathrm{I,II}} T_K^{-1}\n",
    "&=& 2 i\\left[\\begin{array}{rr}\\rho_1 & 0  \\\\ 0 & 0 \\end{array}\\right], \\\\\n",
    "\\operatorname{Disc}_{\\mathrm{I,III}} T_K^{-1}\n",
    "&=& 2 i\\left[\\begin{array}{rr}\\rho_1 & 0  \\\\ 0 & \\rho_2 \\end{array}\\right], \\\\\n",
    "\\operatorname{Disc}_{\\mathrm{I,IV}} T_K^{-1}\n",
    "&=& 2 i\\left[\\begin{array}{rr}0 & 0  \\\\ 0& \\rho_2  \\end{array}\\right].\n",
    "\\end{eqnarray}\n",
    "$$\n",
    "\n",
    "Depending on the centre-of-mass energy, different Riemann sheets connect smoothly to the physical one. Therefore, two cases are studied: one where the resonance mass is above the threshold of the second and first channel, and another where the resonance mass is between the threshold of the first and second channel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "rho = DiagonalMatrix(sp.MatrixSymbol(\"rho\", n, n))\n",
    "Math(aslatex({rho: rho.as_explicit()}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_II = (T_I.inv() + 2 * sp.I * rho).inv()\n",
    "T_III = (T_I.inv() + 2 * sp.I * rho).inv()\n",
    "T_IV = (-T_I.inv() - 2 * sp.I * rho).inv()\n",
    "Math(aslatex([T_II, T_III, T_IV]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "T_II_explicit = T_II.as_explicit()\n",
    "T_II_explicit[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "T_III_explicit = T_III.as_explicit()\n",
    "T_III_explicit[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "T_IV_explicit = T_IV.as_explicit()\n",
    "T_IV_explicit[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "rho_expressions_II = {\n",
    "    **cm_expressions,\n",
    "    rho[0, 0]: PhaseSpaceFactor(s, ma1, mb1),\n",
    "    rho[1, 1]: 0,\n",
    "}\n",
    "rho_expressions_III = {\n",
    "    **cm_expressions,\n",
    "    rho[0, 0]: PhaseSpaceFactor(s, ma1, mb1),\n",
    "    rho[1, 1]: PhaseSpaceFactor(s, ma2, mb2),\n",
    "}\n",
    "rho_expressions_IV = {\n",
    "    **cm_expressions,\n",
    "    rho[0, 0]: 0,\n",
    "    rho[1, 1]: PhaseSpaceFactor(s, ma2, mb2),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "T_II_rho_expr = T_II_explicit.xreplace(rho_expressions_II)\n",
    "T_III_rho_expr = T_III_explicit.xreplace(rho_expressions_III)\n",
    "T_IV_rho_expr = T_IV_explicit.xreplace(rho_expressions_IV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "T_II_rho_expr[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "T_III_rho_expr[0, 0].simplify(doit=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Lineshapes (real axis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "%config InlineBackend.figure_formats = [\"svg\"]\n",
    "\n",
    "T_I_func = sp.lambdify(symbols, T_I_cm_expr[0, 0].doit())\n",
    "T_II_func = sp.lambdify(symbols, T_II_rho_expr[0, 0].doit())\n",
    "T_III_func = sp.lambdify(symbols, T_III_rho_expr[0, 0].doit())\n",
    "T_IV_func = sp.lambdify(symbols, T_IV_rho_expr[0, 0].doit())\n",
    "parameter_defaults1 = {\n",
    "    ma1: 1.0,\n",
    "    mb1: 1.5,\n",
    "    ma2: 1.5,\n",
    "    mb2: 2.0,\n",
    "    m0: 4.0,\n",
    "    g1: 0.7,\n",
    "    g2: 0.7,\n",
    "}\n",
    "parameter_defaults2 = {\n",
    "    **parameter_defaults1,\n",
    "    m0: 3.0,\n",
    "}\n",
    "args1 = eval(str(symbols[1:].xreplace(parameter_defaults1)))\n",
    "args2 = eval(str(symbols[1:].xreplace(parameter_defaults2)))\n",
    "\n",
    "epsilon = 1e-5\n",
    "x = np.linspace(0, 8, num=300)\n",
    "y = np.linspace(epsilon, 1, num=100)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "Zn = X - Y * 1j\n",
    "Zp = X + Y * 1j\n",
    "\n",
    "T1n_res1 = T_I_func(Zn**2, *args1)\n",
    "T1p_res1 = T_I_func(Zp**2, *args1)\n",
    "\n",
    "T2n_res1 = T_II_func(Zn**2, *args1)\n",
    "T2p_res1 = T_II_func(Zp**2, *args1)\n",
    "\n",
    "T3n_res1 = T_III_func(Zn**2, *args1)\n",
    "T3p_res1 = T_III_func(Zp**2, *args1)\n",
    "\n",
    "T4n_res1 = T_IV_func(Zn**2, *args1)\n",
    "T4p_res1 = T_IV_func(Zp**2, *args1)\n",
    "\n",
    "T1n_res2 = T_I_func(Zn**2, *args2)\n",
    "T1p_res2 = T_I_func(Zp**2, *args2)\n",
    "\n",
    "T2n_res2 = T_II_func(Zn**2, *args2)\n",
    "T2p_res2 = T_II_func(Zp**2, *args2)\n",
    "\n",
    "T3n_res2 = T_III_func(Zn**2, *args2)\n",
    "T3p_res2 = T_III_func(Zp**2, *args2)\n",
    "\n",
    "T4n_res2 = T_IV_func(Zn**2, *args2)\n",
    "T4p_res2 = T_IV_func(Zp**2, *args2)\n",
    "\n",
    "fig, axes = plt.subplots(figsize=(11, 6), ncols=4, sharey=True)\n",
    "ax1, ax2, ax3, ax4 = axes.flatten()\n",
    "\n",
    "ax1.plot(x, T1n_res1[0].imag, label=R\"$T_\\mathrm{I}(s-0i)$\")\n",
    "ax1.plot(x, T1p_res1[0].imag, label=R\"$T_\\mathrm{I}(s+0i)$\")\n",
    "ax1.set_title(f\"${sp.latex(rho_cm_expr)}$\")\n",
    "ax1.set_title(R\"$T_\\mathrm{I}$\")\n",
    "\n",
    "ax2.plot(x, T2n_res1[0].imag, label=R\"$T_\\mathrm{II}(s-0i)$\")\n",
    "ax2.plot(x, T2p_res1[0].imag, label=R\"$T_\\mathrm{II}(s+0i)$\")\n",
    "ax2.set_title(R\"$T_\\mathrm{II}$\")\n",
    "\n",
    "ax3.plot(x, T3n_res1[0].imag, label=R\"$T_\\mathrm{III}(s-0i)$\")\n",
    "ax3.plot(x, T3p_res1[0].imag, label=R\"$T_\\mathrm{III}(s+0i)$\")\n",
    "ax3.set_title(R\"$T_\\mathrm{III}$\")\n",
    "\n",
    "ax4.plot(x, T4n_res1[0].imag, label=R\"$T_\\mathrm{III}(s-0i)$\")\n",
    "ax4.plot(x, T4p_res1[0].imag, label=R\"$T_\\mathrm{IV}(s+0i)$\")\n",
    "ax4.set_title(R\"$T_\\mathrm{III}$\")\n",
    "\n",
    "for ax in axes:\n",
    "    ax.legend()\n",
    "    ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n",
    "    ax.set_ylim(-1, +1)\n",
    "ax1.set_ylabel(R\"$\\mathrm{Im}\\,T(s)$ (a.u.)\")\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "### Complex plane (2D)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "tags": []
   },
   "source": [
    "It can be shown that if the resonance mass is above both thresholds the third sheet connects smoothly to the first sheet. If the resonance mass is above the first and below the second threshold the second sheet transitions smoothly into the first sheet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "scroll-input",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "%config InlineBackend.figure_formats = [\"png\"]\n",
    "\n",
    "fig, axes = plt.subplots(figsize=(12, 8), ncols=2, nrows=2, sharey=True)\n",
    "ax1, ax2, ax3, ax4 = axes.flatten()\n",
    "\n",
    "for ax in axes.flatten():\n",
    "    ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n",
    "for ax in axes[:, 0]:\n",
    "    ax.set_ylabel(R\"$\\mathrm{Im}\\,\\sqrt{s}$\")\n",
    "\n",
    "ax1.set_title(\"I and II\")\n",
    "ax2.set_title(\"I and III\")\n",
    "ax3.set_title(\"I and II\")\n",
    "ax4.set_title(\"I and III\")\n",
    "\n",
    "T_max = 2\n",
    "\n",
    "style = dict(vmin=-T_max, vmax=+T_max, cmap=plt.cm.coolwarm)\n",
    "mesh = ax1.pcolormesh(X, Y, T1p_res1.imag, **style)\n",
    "ax1.pcolormesh(X, -Y, T2n_res1.imag, **style)\n",
    "ax2.pcolormesh(X, +Y, T1p_res1.imag, **style)\n",
    "ax2.pcolormesh(X, -Y, T3n_res1.imag, **style)\n",
    "ax3.pcolormesh(X, +Y, T1p_res2.imag, **style)\n",
    "ax3.pcolormesh(X, -Y, T2n_res2.imag, **style)\n",
    "ax4.pcolormesh(X, +Y, T1p_res2.imag, **style)\n",
    "ax4.pcolormesh(X, -Y, T3n_res2.imag, **style)\n",
    "\n",
    "s_thr1 = parameter_defaults1[ma1] + parameter_defaults1[mb1]\n",
    "s_thr2 = parameter_defaults1[ma2] + parameter_defaults1[mb2]\n",
    "linestyle = dict(ls=\"dotted\", lw=1)\n",
    "for ax in axes.flatten():\n",
    "    ax.axhline(0, c=\"black\", **linestyle)\n",
    "    ax.axvline(s_thr1, c=\"C0\", **linestyle, label=R\"$\\sqrt{s_\\mathrm{thr1}}$\")\n",
    "    ax.axvline(s_thr2, c=\"C1\", **linestyle, label=R\"$\\sqrt{s_\\mathrm{thr2}}$\")\n",
    "linestyle = dict(c=\"r\", ls=\"dotted\", label=R\"$m_\\mathrm{res}$\")\n",
    "for ax in axes[0]:\n",
    "    ax.axvline(parameter_defaults1[m0], **linestyle)\n",
    "for ax in axes[1]:\n",
    "    ax.axvline(parameter_defaults2[m0], **linestyle)\n",
    "ax2.legend()\n",
    "\n",
    "fig.text(0.5, 0.93, R\"$s_{thr1}<s_{thr2}<m_{res}$\", ha=\"center\", fontsize=18)\n",
    "fig.text(0.5, 0.46, R\"$s_{thr1}<m_{res}<s_{thr2}$\", ha=\"center\", fontsize=18)\n",
    "fig.subplots_adjust(wspace=1)\n",
    "cax = fig.add_axes([0.92, 0.15, 0.02, 0.7])\n",
    "cbar = fig.colorbar(mesh, cax=cax)\n",
    "cbar.ax.set_title(R\"$\\mathrm{Im} T(s)$\")\n",
    "\n",
    "fig.tight_layout(rect=[0, 0.03, 0.9, 0.95])\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Riemann sheets (3D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Define plot style"
    },
    "tags": [
     "scroll-input",
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "def sty(sheet_name: str) -> dict:\n",
    "    sheet_color = sheet_colors[sheet_name]\n",
    "    n_lines = 16\n",
    "    return dict(\n",
    "        cmin=-vmax,\n",
    "        cmax=+vmax,\n",
    "        colorscale=[[0, \"rgb(0, 0, 0)\"], [1, sheet_color]],\n",
    "        contours=dict(\n",
    "            x=dict(\n",
    "                show=True,\n",
    "                start=x.min(),\n",
    "                end=x.max(),\n",
    "                size=(x.max() - x.min()) / n_lines,\n",
    "                color=\"black\",\n",
    "            ),\n",
    "            y=dict(\n",
    "                show=True,\n",
    "                start=-y.max(),\n",
    "                end=+y.max(),\n",
    "                size=(y.max() - y.min()) / (n_lines // 2),\n",
    "                color=\"black\",\n",
    "            ),\n",
    "        ),\n",
    "        name=sheet_name,\n",
    "        opacity=0.4,\n",
    "        showscale=False,\n",
    "    )\n",
    "\n",
    "\n",
    "vmax = 2.0\n",
    "project = np.imag\n",
    "sheet_colors = {\n",
    "    \"T1 (physical)\": \"blue\",\n",
    "    \"T2 (unphysical)\": \"red\",\n",
    "    \"T3 (unphysical)\": \"green\",\n",
    "    \"T4 (unphysical)\": \"yellow\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "Sp_I_res1 = go.Surface(x=X, y=+Y, z=T1p_res1.imag, **sty(\"T1 (physical)\"))\n",
    "Sn_II_res1 = go.Surface(x=X, y=-Y, z=T2n_res1.imag, **sty(\"T2 (unphysical)\"))\n",
    "Sn_III_res1 = go.Surface(x=X, y=-Y, z=T3n_res1.imag, **sty(\"T3 (unphysical)\"))\n",
    "\n",
    "Sp_I_res2 = go.Surface(x=X, y=+Y, z=T1p_res2.imag, **sty(\"T1 (physical)\"))\n",
    "Sn_II_res2 = go.Surface(x=X, y=-Y, z=T2n_res2.imag, **sty(\"T2 (unphysical)\"))\n",
    "Sn_III_res2 = go.Surface(x=X, y=-Y, z=T3n_res2.imag, **sty(\"T3 (unphysical)\"))\n",
    "\n",
    "thr1_filter = x >= s_thr1\n",
    "thr2_filter = x >= s_thr2\n",
    "\n",
    "line_kwargs = dict(\n",
    "    line=dict(color=\"yellow\", width=8),\n",
    "    mode=\"lines\",\n",
    "    name=\"Lineshape\",\n",
    ")\n",
    "lineshape_res1_z = project(T1p_res1[0])\n",
    "lineshape_res2_z = project(T1p_res2[0])\n",
    "lineshape_res1 = go.Scatter3d(\n",
    "    x=x[thr1_filter],\n",
    "    y=np.zeros(thr1_filter.shape),\n",
    "    z=lineshape_res1_z[thr1_filter],\n",
    "    **line_kwargs,\n",
    ")\n",
    "lineshape_res2 = go.Scatter3d(\n",
    "    x=x[thr1_filter],\n",
    "    y=np.zeros(thr1_filter.shape),\n",
    "    z=lineshape_res2_z[thr1_filter],\n",
    "    **line_kwargs,\n",
    ")\n",
    "\n",
    "point_kwargs = dict(\n",
    "    hoverinfo=\"text\",\n",
    "    marker=dict(color=DEFAULT_PLOTLY_COLORS[:2], size=6),\n",
    "    mode=\"markers\",\n",
    "    text=[\"threshold 1\", \"threshold 2\"],\n",
    ")\n",
    "thr_points_res1 = go.Scatter3d(\n",
    "    x=[s_thr1, s_thr2],\n",
    "    y=[0, 0],\n",
    "    z=[lineshape_res1_z[thr1_filter][0], lineshape_res1_z[thr2_filter][0]],\n",
    "    **point_kwargs,\n",
    ")\n",
    "thr_points_res2 = go.Scatter3d(\n",
    "    x=[s_thr1, s_thr2],\n",
    "    y=[0, 0],\n",
    "    z=[lineshape_res2_z[thr1_filter][0], lineshape_res2_z[thr2_filter][0]],\n",
    "    **point_kwargs,\n",
    ")\n",
    "\n",
    "plotly_fig = make_subplots(\n",
    "    rows=2,\n",
    "    cols=2,\n",
    "    horizontal_spacing=0.01,\n",
    "    vertical_spacing=0.05,\n",
    "    specs=[\n",
    "        [{\"type\": \"surface\"}, {\"type\": \"surface\"}],\n",
    "        [{\"type\": \"surface\"}, {\"type\": \"surface\"}],\n",
    "    ],\n",
    "    subplot_titles=[\n",
    "        \"thr₁ < thr₂ < mᵣ\",\n",
    "        \"thr₁ < mᵣ < thr₂\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# thr₁ < thr₂ < mᵣ\n",
    "selector = dict(col=1, row=1)\n",
    "plotly_fig.add_trace(Sp_I_res1, **selector)\n",
    "plotly_fig.add_trace(Sn_III_res1, **selector)\n",
    "plotly_fig.add_trace(lineshape_res1, **selector)\n",
    "plotly_fig.add_trace(thr_points_res1, **selector)\n",
    "selector = dict(col=1, row=2)\n",
    "plotly_fig.add_trace(Sp_I_res1, **selector)\n",
    "plotly_fig.add_trace(Sn_II_res1, **selector)\n",
    "plotly_fig.add_trace(lineshape_res1, **selector)\n",
    "plotly_fig.add_trace(thr_points_res1, **selector)\n",
    "\n",
    "# thr₁ < mᵣ < thr₂\n",
    "selector = dict(col=2, row=1)\n",
    "plotly_fig.add_trace(Sp_I_res2, **selector)\n",
    "plotly_fig.add_trace(Sn_II_res2, **selector)\n",
    "plotly_fig.add_trace(lineshape_res2, **selector)\n",
    "plotly_fig.add_trace(thr_points_res2, **selector)\n",
    "selector = dict(col=2, row=2)\n",
    "plotly_fig.add_trace(Sp_I_res2, **selector)\n",
    "plotly_fig.add_trace(Sn_III_res2, **selector)\n",
    "plotly_fig.add_trace(lineshape_res2, **selector)\n",
    "plotly_fig.add_trace(thr_points_res2, **selector)\n",
    "\n",
    "plotly_fig.update_layout(\n",
    "    height=600,\n",
    "    margin=dict(l=0, r=0, t=20, b=0),\n",
    "    showlegend=False,\n",
    ")\n",
    "\n",
    "plotly_fig.update_scenes(\n",
    "    camera_center=dict(z=-0.1),\n",
    "    camera_eye=dict(x=1.4, y=1.4, z=1.4),\n",
    "    xaxis_range=(2.0, 5.0),\n",
    "    xaxis_title_text=\"Re √s\",\n",
    "    yaxis_title_text=\"Im √s\",\n",
    "    zaxis_title_text=\"Im T(s)\",\n",
    "    zaxis_range=[-vmax, +vmax],\n",
    ")\n",
    "plotly_fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complex plane widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "editable": true,
    "tags": [
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "sliders = dict(\n",
    "    g01=w.FloatSlider(\n",
    "        description=R\"$g^{0}_1$\",\n",
    "        value=0.5,\n",
    "        min=-2,\n",
    "        max=+2,\n",
    "    ),\n",
    "    g02=w.FloatSlider(\n",
    "        description=\"$g^{0}_2$\",\n",
    "        value=0.5,\n",
    "        min=-2,\n",
    "        max=+2,\n",
    "    ),\n",
    "    m0=w.FloatSlider(\n",
    "        description=\"$m_0$\",\n",
    "        value=4,\n",
    "        min=0,\n",
    "        max=+6,\n",
    "    ),\n",
    "    T_max=w.FloatSlider(\n",
    "        description=R\"$T_\\mathrm{max}$\",\n",
    "        value=1.0,\n",
    "        min=0.1,\n",
    "        max=6.0,\n",
    "        step=0.1,\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "fig, axes = plt.subplots(\n",
    "    figsize=(11, 6),\n",
    "    ncols=2,\n",
    "    nrows=2,\n",
    "    sharex=True,\n",
    "    gridspec_kw={\"height_ratios\": [1, 2]},\n",
    ")\n",
    "fig.canvas.toolbar_visible = False\n",
    "fig.canvas.header_visible = False\n",
    "fig.canvas.footer_visible = False\n",
    "ax1d1, ax1d2, ax2d1, ax2d2 = axes.flatten()\n",
    "\n",
    "for ax in axes[1]:\n",
    "    ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n",
    "ax1d1.set_ylabel(\"Intensity (a.u.)\")\n",
    "ax2d1.set_ylabel(R\"$\\mathrm{Im}\\,\\sqrt{s}$\")\n",
    "\n",
    "ax1d1.set_title(\"I and II\")\n",
    "ax1d2.set_title(\"I and III\")\n",
    "\n",
    "R_color = \"C4\"\n",
    "T1_color = \"C0\"\n",
    "T2_color = \"C3\"\n",
    "T3_color = \"C2\"\n",
    "\n",
    "\n",
    "LINES = None\n",
    "MESH = None\n",
    "style = dict(cmap=plt.cm.coolwarm)\n",
    "\n",
    "\n",
    "def plot(m0, g01, g02, T_max):\n",
    "    global LINES, MESH\n",
    "    local_args = args1[:-3] + (m0, g01, g02)\n",
    "    T1p_res1 = T_I_func(Zp**2, *local_args)\n",
    "    T2n_res1 = T_II_func(Zn**2, *local_args)\n",
    "    T3n_res1 = T_III_func(Zn**2, *local_args)\n",
    "    T1y = np.abs(T1p_res1[0]) ** 2\n",
    "    T2y = np.abs(T2n_res1[0]) ** 2\n",
    "    T3y = np.abs(T3n_res1[0]) ** 2\n",
    "    if MESH is None and LINES is None:\n",
    "        LINES = [\n",
    "            ax1d1.axvline(m0, c=R_color, ls=\"dashed\"),\n",
    "            ax1d2.axvline(m0, c=R_color, ls=\"dashed\"),\n",
    "            ax2d1.axvline(m0, c=R_color, ls=\"dashed\", label=R\"$m_\\mathrm{res}$\"),\n",
    "            ax2d2.axvline(m0, c=R_color, ls=\"dashed\", label=R\"$m_\\mathrm{res}$\"),\n",
    "            ax1d1.plot(x, T1y, c=T1_color, label=R\"$\\left|T_\\mathrm{I}\\right|^2$\")[0],\n",
    "            ax1d1.plot(\n",
    "                x, T2y, c=T2_color, label=R\"$\\left|T_\\mathrm{II}\\right|^2$\", ls=\"dotted\"\n",
    "            )[0],\n",
    "            ax1d2.plot(x, T1y, c=T1_color, label=R\"$\\left|T_\\mathrm{I}\\right|^2$\")[0],\n",
    "            ax1d2.plot(\n",
    "                x,\n",
    "                T3y,\n",
    "                c=T3_color,\n",
    "                label=R\"$\\left|T_\\mathrm{III}\\right|^2$\",\n",
    "                ls=\"dotted\",\n",
    "            )[0],\n",
    "        ]\n",
    "        MESH = [\n",
    "            ax2d1.pcolormesh(X, Y, T1p_res1.imag, **style),\n",
    "            ax2d1.pcolormesh(X, -Y, T2n_res1.imag, **style),\n",
    "            ax2d2.pcolormesh(X, +Y, T1p_res1.imag, **style),\n",
    "            ax2d2.pcolormesh(X, -Y, T3n_res1.imag, **style),\n",
    "        ]\n",
    "    else:\n",
    "        MESH[0].set_array(T1p_res1.imag)\n",
    "        MESH[1].set_array(T2n_res1.imag)\n",
    "        MESH[2].set_array(T1p_res1.imag)\n",
    "        MESH[3].set_array(T3n_res1.imag)\n",
    "        LINES[0].set_xdata(m0)\n",
    "        LINES[1].set_xdata(m0)\n",
    "        LINES[2].set_xdata(m0)\n",
    "        LINES[3].set_xdata(m0)\n",
    "        LINES[4].set_ydata(T1y)\n",
    "        LINES[5].set_ydata(T2y)\n",
    "        LINES[6].set_ydata(T1y)\n",
    "        LINES[7].set_ydata(T3y)\n",
    "    for mesh in MESH:\n",
    "        mesh.set_clim(-T_max, +T_max)\n",
    "    for ax in axes[0]:\n",
    "        ax.set_ylim(0, max(T1y) * 1.05)\n",
    "    fig.canvas.draw()\n",
    "\n",
    "\n",
    "for ax in axes[:, 1]:\n",
    "    ax.set_yticks([])\n",
    "for ax in axes[1]:\n",
    "    ax.axhline(0, c=\"black\", ls=\"dotted\", lw=1)\n",
    "for ax in axes[0]:\n",
    "    ax.axvline(s_thr1, c=\"C0\", ls=\"dotted\", lw=1)\n",
    "    ax.axvline(s_thr2, c=\"C1\", ls=\"dotted\", lw=1)\n",
    "for ax in axes[1]:\n",
    "    ax.axvline(s_thr1, c=\"C0\", label=R\"$\\sqrt{s_\\mathrm{thr1}}$\", ls=\"dotted\", lw=1)\n",
    "    ax.axvline(s_thr2, c=\"C1\", label=R\"$\\sqrt{s_\\mathrm{thr2}}$\", ls=\"dotted\", lw=1)\n",
    "ax2d1.text(0.5, +0.7, R\"$T_\\mathrm{I}$\", color=T1_color, size=20)\n",
    "ax2d1.text(0.5, -0.75, R\"$T_\\mathrm{II}$\", color=T2_color, size=20)\n",
    "ax2d2.text(0.5, +0.7, R\"$T_\\mathrm{I}$\", color=T1_color, size=20)\n",
    "ax2d2.text(0.5, -0.75, R\"$T_\\mathrm{III}$\", color=T3_color, size=20)\n",
    "\n",
    "output = w.interactive_output(plot, controls=sliders)\n",
    "UI = w.VBox(list(sliders.values()))\n",
    "fig.tight_layout()\n",
    "ax1d1.legend()\n",
    "ax1d2.legend()\n",
    "ax2d2.legend()\n",
    "display(output, UI)"
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}