FEAT: implement different algorithms for NumericalIntegral

Author: redeboer

docs/usage/sympy.ipynb

@@ -399,6 +399,137 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that the choice of {attr}`.NumericalIntegral.algorithm` is important. For vectorized integrals with large input arrays, it is better to use [JAX](https://docs.jax.dev) as a numerical backend and any of the [`quadax`](http://quadax.readthedocs.io) algorithms as the integration method.\n",
+    "\n",
+    "Below, we see a comparison between the {func}`quadax.romberg` and {func}`quadax.quadgk` algorithms for a contour integral in the complex plane. The {func}`quadax.quadgk` algorithm is much faster than the {func}`quadax.romberg` algorithm, but is less accurate when $a$ is close to the contour."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "a, z = sp.symbols(\"a z\")\n",
+    "f = z**3 / (z - a)\n",
+    "f"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "r, phi = sp.symbols(\"r phi\")\n",
+    "f_C = f.subs(z, r * sp.exp(sp.I * phi))\n",
+    "quadgk_expr = NumericalIntegral(f_C, (phi, 0, 2 * sp.pi), algorithm=\"quadax.quadgk\")\n",
+    "romberg_expr = NumericalIntegral(f_C, (phi, 0, 2 * sp.pi), algorithm=\"quadax.romberg\")\n",
+    "romberg_expr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "args = (a, r)\n",
+    "quadgk_func = sp.lambdify(args, quadgk_expr, modules=\"jax\")\n",
+    "romberg_func = sp.lambdify(args, romberg_expr, modules=\"jax\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "\n",
+    "import jax.numpy as jnp\n",
+    "\n",
+    "x_max = 4\n",
+    "X, Y = jnp.meshgrid(\n",
+    "    jnp.linspace(-x_max, +x_max, num=300),\n",
+    "    jnp.linspace(-x_max, +x_max, num=300),\n",
+    ")\n",
+    "start = time.perf_counter()\n",
+    "Z = quadgk_func(a=X + Y * 1j, r=3).block_until_ready()\n",
+    "end = time.perf_counter()\n",
+    "print(f\"Computation took {end - start:.2f} seconds for {X.shape[0]}x{X.shape[1]} grid\")\n",
+    "\n",
+    "z_max = 30\n",
+    "fig, axes = plt.subplots(ncols=2, figsize=(8.5, 5), sharey=True)\n",
+    "fig.suptitle(\"Integral $I_a$ computed with quadax.quadgk\", y=0.95)\n",
+    "ax_real, ax_imag = axes\n",
+    "ax_real.set_title(\"Real part of $I_a$\")\n",
+    "ax_imag.set_title(\"Imaginary part of $I_a$\")\n",
+    "ax_real.set_xlabel(\"Re $a$\")\n",
+    "ax_real.set_ylabel(\"Im $a$\")\n",
+    "ax_imag.set_xlabel(\"Re $a$\")\n",
+    "for ax, Z_proj in zip(axes, [Z.real, Z.imag], strict=True):\n",
+    "    ax.pcolormesh(\n",
+    "        X, Y, Z_proj, cmap=\"RdBu_r\", rasterized=True, vmin=-z_max, vmax=+z_max\n",
+    "    )\n",
+    "fig.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "X, Y = jnp.meshgrid(\n",
+    "    jnp.linspace(-x_max, +x_max, num=50),\n",
+    "    jnp.linspace(-x_max, +x_max, num=50),\n",
+    ")\n",
+    "start = time.perf_counter()\n",
+    "Z = romberg_func(a=X + Y * 1j, r=3).block_until_ready()\n",
+    "end = time.perf_counter()\n",
+    "print(f\"Computation took {end - start:.2f} seconds for {X.shape[0]}x{X.shape[1]} grid\")\n",
+    "\n",
+    "fig, axes = plt.subplots(ncols=2, figsize=(8.5, 5), sharey=True)\n",
+    "fig.suptitle(\"Integral $I_a$ computed with quadax.romberg\", y=0.95)\n",
+    "ax_real, ax_imag = axes\n",
+    "ax_real.set_title(\"Real part\")\n",
+    "ax_imag.set_title(\"Imaginary part\")\n",
+    "ax_real.set_xlabel(\"Re $a$\")\n",
+    "ax_real.set_ylabel(\"Im $a$\")\n",
+    "ax_imag.set_xlabel(\"Re $a$\")\n",
+    "for ax, z in zip(axes, [Z.real, Z.imag], strict=True):\n",
+    "    ax.pcolormesh(X, Y, z, cmap=\"RdBu_r\", rasterized=True, vmin=-z_max, vmax=+z_max)\n",
+    "fig.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    ":::{seealso}\n",
+    "The performance of different integration algorithms for computing the dispersion integral in different scenarios is investigated in [this notebook](./dynamics/integration-algorithms.ipynb).\n",
+    ":::"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

pyproject.toml

@@ -118,6 +118,7 @@ test = [
     "pytest-cov",
     "pytest-profiling",
     "pytest-xdist",
+    "quadax",
 ]
 types = [
     "IPython",
@@ -365,6 +366,8 @@ filterwarnings = [
     "ignore:datetime.datetime.utcfromtimestamp\\(\\) is deprecated and scheduled for removal in a future version.*:DeprecationWarning",
     "ignore:unclosed .*:ResourceWarning",
     'ignore:Widget\..* is deprecated\.:DeprecationWarning',
+    'ignore:atleast_1d requires ndarray or scalar arguments.*:DeprecationWarning',
+    'ignore:jax\.core\.(un)?mapped_aval is deprecated\.:DeprecationWarning',
 ]
 markers = ["slow: marks tests as slow (select with '-m slow')"]
 minversion = "9.0"

src/ampform/dynamics/phasespace.py

@@ -298,12 +298,15 @@ class ChewMandelstamIntegral(sp.Expr):
     meson_radius: Any = 1
     name: str | None = argument(default=None, sympify=False)
     algorithm: tuple[str, str] | None = argument(
-        default=None, sympify=False, kw_only=True
+        default=None, kw_only=True, sympify=False
     )
+    """See :attr:`.NumericalIntegral.algorithm`."""
     configuration: dict[str, Any] | None = argument(
-        default=None, sympify=False, kw_only=True
+        default=None, kw_only=True, sympify=False
     )
-    dummify: bool = argument(default=True, sympify=False, kw_only=True)
+    """See :attr:`.NumericalIntegral.configuration`."""
+    dummify: bool = argument(default=True, kw_only=True, sympify=False)
+    """Whether to dummify the integration variable. See :attr:`.NumericalIntegral.dummify`."""
 
     def evaluate(self) -> sp.Expr:
         s, m1, m2, L, s_prime, epsilon, meson_radius, *_ = self.args  # noqa: N806

src/ampform/sympy/__init__.py

@@ -15,7 +15,6 @@
 import itertools
 import re
 import sys
-import warnings
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, cast
 
@@ -373,14 +372,35 @@ class NumericalIntegral(sp.Integral):
     """
 
     function: sp.Expr
+    """Integrand of the integral."""
     limits: tuple[sp.Symbol, sp.Basic, sp.Basic]
-    algorithm: tuple[str, str] | None = argument(
-        default=None, sympify=False, kw_only=True
-    )
+    """Integration variable and its limits (can be `~sympy.core.numbers.Infinity`)."""
+    algorithm: str | None = argument(default=None, kw_only=True, sympify=False)
+    """Name of the numerical integration algorithm to use when lambdifying this integral.
+
+    The algorithm should be in the format :code:`module.function`, for instance
+    :func:`scipy.integrate.quad_vec` or :func:`quadax.quadgk`. By default, the algorithm
+    is :func:`quadax.romberg` when lambdifying to JAX and
+    :func:`scipy.integrate.quad_vec` when lambdifying to NumPy.
+    """
     configuration: dict[str, Any] | None = argument(
         default=None, sympify=False, kw_only=True
     )
+    """Keyword arguments for the numerical integration algorithm.
+
+    For example, for :func:`scipy.integrate.quad_vec`, one can set the relative
+    tolerance with :code:`configuration={'epsrel': 1e-5}`.
+    """
     dummify: bool = argument(default=True, sympify=False, kw_only=True)
+    """Replace the integration variable with a dummy symbol before lambdification.
+
+    The integrand expression is lambdified to a :code:`lambda` function. Therefore, when
+    the integrand expresssion contains the integration variable in a non-trivial way,
+    and the expression is lambdified using common sub-expressions, it is better to
+    replace it with a unique `~sympy.core.symbol.Dummy` symbol that does not appear
+    anywhere else in the expression tree, so that is not pulled out of the
+    :code:`lambda` function.
+    """
 
     @override
     def doit(self, **hints):
@@ -391,54 +411,82 @@ def doit(self, **hints):
         }
         return self.func(*args, **kwargs)
 
+    @override
+    def _jaxcode(self, printer, *args) -> str:  # ty:ignore[invalid-explicit-override]
+        algorithm = self.algorithm or "quadax.romberg"
+        if algorithm.startswith("quadax"):
+            return self.__to_quadax_like(printer, algorithm)
+        return self.__to_scipy_like(printer, algorithm)
+
     @override
     def _numpycode(self, printer, *args) -> str:  # ty:ignore[invalid-explicit-override]
-        module, algorithm = self.algorithm or ("scipy.integrate", "quad_vec")
-        if module.startswith("scipy"):
-            _warn_if_scipy_not_installed()
+        algorithm = self.algorithm or "scipy.integrate.quad_vec"
+        if algorithm.startswith("quadax"):
+            return self.__to_quadax_like(printer, algorithm)
+        return self.__to_scipy_like(printer, algorithm)
+
+    def __to_quadax_like(self, printer, algorithm: str) -> str:
+        """https://quadax.readthedocs.io."""
+        integrate, integrand, x, a, b = self.__prepare_components(printer, algorithm)
+        src = _generate_function_call(
+            integrate,
+            fun=f"lambda {x}: {integrand}",
+            interval=f"({a}, {b})",
+            **self.configuration or {},
+        )
+        return f"{src}[0]"
+
+    def __to_scipy_like(self, printer, algorithm: str) -> str:
+        """https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad_vec.html."""
+        integrate, integrand, x, a, b = self.__prepare_components(printer, algorithm)
+        kwargs = self.configuration or {}
+        src = _generate_function_call(
+            integrate, f"lambda {x}: {integrand}", a, b, **kwargs
+        )
+        return f"{src}[0]"
+
+    def __prepare_components(
+        self, printer, algorithm: str
+    ) -> tuple[str, str, str, str, str]:
         integration_vars, limits = _unpack_integral_limits(self)
         if len(limits) != 1 or len(integration_vars) != 1:
             msg = f"Cannot handle {len(limits)}-dimensional integrals"
             raise ValueError(msg)
         x = integration_vars[0]
         a, b = limits[0]
-        expr = self.args[0]
+        integrand = self.function
         if self.dummify:
             dummy = sp.Dummy()
-            expr = expr.xreplace({x: dummy})
+            integrand = integrand.xreplace({x: dummy})
             x = dummy
-        integrate_numerically = algorithm
-        printer.module_imports[module].add(integrate_numerically)
-        src = _generate_function_call(
-            # https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad_vec.html
-            integrate_numerically,
-            f"lambda {printer._print(x)}: {printer._print(expr)}",
+        parts = algorithm.split(".")
+        if len(parts) < 2:  # noqa: PLR2004
+            msg = f"Algorithm should be in format 'module.function', got '{algorithm}'"
+            raise ValueError(msg)
+        module_name = ".".join(parts[:-1])
+        algorithm_name = parts[-1]
+        printer.module_imports[module_name].add(algorithm_name)
+        return (
+            algorithm_name,
+            printer._print(integrand),
+            printer._print(x),
             printer._print(a),
             printer._print(b),
-            **self.configuration or {},
         )
-        return f"{src}[0]"
 
 
-def _generate_function_call(func_name: str, *args, **kwargs) -> str:
+def _generate_function_call(func_name: str, /, *args, **kwargs) -> str:
     """Generate a function call string with the given function name, arguments, and keyword arguments.
 
     >>> _generate_function_call("quad_vec", "f", 0, 1, epsabs=1e-5)
     'quad_vec(f, 0, 1, epsabs=1e-05)'
+    >>> _generate_function_call("quadgk", fun="lambda x: x**2", interval=(0, 1))
+    'quadgk(fun=lambda x: x**2, interval=(0, 1))'
     """
-    src = f"{func_name}({', '.join(map(str, args))}"
-    for key, value in kwargs.items():
-        src += f", {key}={value}"
+    src = f"{func_name}("
+    src += ", ".join(map(str, args))
+    if args:
+        src += ", "
+    src += ", ".join(f"{key}={value}" for key, value in kwargs.items())
     src += ")"
     return src
-
-
-def _warn_if_scipy_not_installed() -> None:
-    try:
-        import scipy  # noqa: F401, PLC0415
-    except ImportError:
-        warnings.warn(
-            "Scipy is not installed. Install with 'pip install scipy' or with 'pip"
-            " install ampform[scipy]'",
-            stacklevel=1,
-        )

tests/sympy/test_integral.py

@@ -6,17 +6,40 @@
 
 
 class TestNumericalIntegral:
-    @pytest.mark.parametrize("call_doit", [True, False])
-    @pytest.mark.parametrize("configuration", [{}, {"limit": 10}, {"limit": 100}])
+    @pytest.mark.parametrize(
+        ("backend", "algorithm", "configuration"),
+        [
+            ("jax", "quadax.quadgk", None),
+            ("jax", "quadax.romberg", None),
+            ("numpy", "scipy.integrate.quad_vec", None),
+            ("numpy", "scipy.integrate.quad_vec", {"limit": 10}),
+            ("numpy", None, None),
+        ],
+    )
+    @pytest.mark.parametrize("call_doit", [False, True])
+    @pytest.mark.parametrize("dummify", [False, True])
     def test_real_value_function(
-        self, call_doit: bool, configuration: dict[str, int | None]
+        self,
+        algorithm: str | None,
+        backend: str,
+        call_doit: bool,
+        configuration: dict[str, int | None],
+        dummify: bool,
     ):
         x = sp.symbols("x")
-        integral_expr = NumericalIntegral(x**2, (x, 1, 3), configuration=configuration)
+        integral_expr = NumericalIntegral(
+            x**2,
+            (x, 1, 3),
+            algorithm=algorithm,
+            configuration=configuration,
+            dummify=dummify,
+        )
         if call_doit:
             integral_expr = integral_expr.doit()
+        assert integral_expr.algorithm == algorithm
         assert integral_expr.configuration == configuration
-        func = sp.lambdify(args=[], expr=integral_expr)
+        assert integral_expr.dummify is dummify
+        func = sp.lambdify([], integral_expr, backend)
         assert func() == 26 / 3  # noqa: RUF069
 
     @pytest.mark.parametrize(