ENH: use quadax.quadgk in Riemann sheet notebook

* FIX: set correct LaTeX syntax in equation

Author: redeboer

docs/usage/dynamics/riemann-sheets.ipynb

@@ -40,8 +40,8 @@
    "source": [
     "import warnings\n",
     "\n",
+    "import jax.numpy as jnp\n",
     "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
     "import sympy as sp\n",
     "from IPython.display import Math\n",
     "from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
@@ -166,7 +166,11 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The general form of the Chew--Mandelstam dispersion integral is the following:"
+    "The general form of the Chew--Mandelstam dispersion integral is the following:\n",
+    "\n",
+    ":::{margin}\n",
+    "We set the {attr}`.ChewMandelstamIntegral.algorithm` to {func}`quadax.quadgk` for fast numerical evaluation over the complex plane. For higher precision close the the physical axis (at the cost of performance), choose {func}`quadax.romberg` instead. See also {doc}`integration-algorithms`.\n",
+    "::::"
    ]
   },
   {
@@ -176,7 +180,9 @@
    "outputs": [],
    "source": [
     "s, m1, m2, L = sp.symbols(\"s, m1, m2, L\", nonnegative=True)\n",
-    "integral_expr = ChewMandelstamIntegral(s, m1, m2, L)\n",
+    "integral_expr = ChewMandelstamIntegral(\n",
+    "    s, m1, m2, L, algorithm=\"quadax.quadgk\", configuration={\"order\": 61}\n",
+    ")\n",
     "integral_expr.doit(deep=False)"
    ]
   },
@@ -274,14 +280,14 @@
     "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. For the 2 channel one gets:\n",
     "\n",
     "$$\n",
-    "\\begin{eqnarray}\n",
+    "\\begin{aligned}\n",
     "\\operatorname{Disc}_{\\mathrm{I,II}} \\varSigma_\\ell(s)\n",
-    "&=& 2 i\\left[\\begin{array}{rr}\\rho_1 n_{\\ell,1}^2(s) & 0  \\\\ 0 & 0 \\end{array}\\right], \\\\\n",
+    "\\; &= \\; 2 i\\begin{pmatrix}\\rho_1(s) \\, n_{\\ell,1}^2(s) & 0  \\\\ 0 & 0 \\end{pmatrix}, \\\\\n",
     "\\operatorname{Disc}_{\\mathrm{I,III}} \\varSigma_\\ell(s)\n",
-    "&=& 2 i\\left[\\begin{array}{rr}\\rho_1 n_{\\ell,1}^2(s) & 0  \\\\ 0 & \\rho_2 n_{\\ell,2}^2(s) \\end{array}\\right], \\\\\n",
+    "\\; &= \\; 2 i\\begin{pmatrix}\\rho_1(s) \\, n_{\\ell,1}^2(s) & 0  \\\\ 0 & \\rho_2(s) \\, n_{\\ell,2}^2(s) \\end{pmatrix}, \\\\\n",
     "\\operatorname{Disc}_{\\mathrm{I,IV}} \\varSigma_\\ell(s)\n",
-    "&=& 2 i\\left[\\begin{array}{rr}0 & 0  \\\\ 0& \\rho_2  n_{\\ell,2}^2(s) \\end{array}\\right].\n",
-    "\\end{eqnarray}\n",
+    "\\; &= \\; 2 i\\begin{pmatrix}0 & 0  \\\\ 0& \\rho_2(s) \\,  n_{\\ell,2}^2(s) \\end{pmatrix}.\n",
+    "\\end{aligned}\n",
     "$$"
    ]
   },
@@ -431,6 +437,7 @@
     "integral_p_wave_func = sp.lambdify(\n",
     "    [s, m1, m2, integral_expr.epsilon],\n",
     "    integral_expr.subs(L, L_val).doit(),\n",
+    "    modules=\"jax\",\n",
     ")"
    ]
   },
@@ -477,9 +484,9 @@
    },
    "outputs": [],
    "source": [
-    "x = np.linspace(0, 6, num=150)\n",
-    "y = np.linspace(1e-3, 1, num=40)\n",
-    "X, Y = np.meshgrid(x, y)\n",
+    "x = jnp.linspace(0, 6, num=200)\n",
+    "y = jnp.linspace(5e-3, 1, num=60)\n",
+    "X, Y = jnp.meshgrid(x, y)\n",
     "Zn = X - Y * 1j\n",
     "Zp = X + Y * 1j"
    ]
@@ -581,7 +588,7 @@
    },
    "outputs": [],
    "source": [
-    "x = np.linspace(2.2, 8, num=100)\n",
+    "x = jnp.linspace(2.2, 8, num=100)\n",
     "ε_offset = 1e-2\n",
     "cn = x - ε_offset * 1j\n",
     "cp = x + ε_offset * 1j\n",
@@ -619,15 +626,6 @@
     "plt.show()"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    ":::{caution}\n",
-    "The integral does not appear to be numerically stable (see [ComPWA/ampform#487](https://github.com/ComPWA/ampform/issues/487)), as a significant imaginary offset from the real axis ($\\epsilon = 10^{-2}i$) is required to get rid of the instability fluctuations.\n",
-    ":::"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -655,7 +653,7 @@
    },
    "outputs": [],
    "source": [
-    "y = np.linspace(1e-8, 1, num=100)\n",
+    "y = jnp.linspace(1e-4, 1, num=80)\n",
     "for s_value in (3.1, 3.7):\n",
     "    Jn = s_value - y * 1j\n",
     "    Jp = s_value + y * 1j\n",