Lambdifying a symbolic integral#

Hide code cell source
import inspect

import black
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from import aslatex
from ampform.sympy import unevaluated
from IPython.display import Markdown, Math
from scipy.integrate import quad, quad_vec
from sympy.printing.pycode import _unpack_integral_limits

Numerical integration#

SciPy’s quad() function#

SciPy’s scipy.integrate.quad() cannot integrate complex-valued functions:

def integrand(x):
    return x * (x + 1j)

quad(integrand, 0.0, 2.0)
TypeError                                 Traceback (most recent call last)
Cell In[3], line 5
      1 def integrand(x):
      2     return x * (x + 1j)
----> 5 quad(integrand, 0.0, 2.0)

File /opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/scipy/integrate/, in quad(func, a, b, args, full_output, epsabs, epsrel, limit, points, weight, wvar, wopts, maxp1, limlst, complex_func)
    461     return retval
    463 if weight is None:
--> 464     retval = _quad(func, a, b, args, full_output, epsabs, epsrel, limit,
    465                    points)
    466 else:
    467     if points is not None:

File /opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/scipy/integrate/, in _quad(func, a, b, args, full_output, epsabs, epsrel, limit, points)
    609 if points is None:
    610     if infbounds == 0:
--> 611         return _quadpack._qagse(func,a,b,args,full_output,epsabs,epsrel,limit)
    612     else:
    613         return _quadpack._qagie(func, bound, infbounds, args, full_output, 
    614                                 epsabs, epsrel, limit)

TypeError: must be real number, not complex

A proposed solution is to wrap the quad() function in a special integrate function that integrates the real and imaginary part of a function separately:

def complex_integrate(func, a, b, **quad_kwargs):
    def real_func(x):
        return func(x).real

    def imag_func(x):
        return func(x).imag

    real_integral, real_integral_err = quad(real_func, a, b, **quad_kwargs)
    imag_integral, imag_integral_err = quad(imag_func, a, b, **quad_kwargs)
    return (
        real_integral + 1j * imag_integral,
        real_integral_err**2 + 1j * imag_integral_err,

complex_integrate(integrand, 0.0, 2.0)
((2.666666666666667+2j), (8.765121169122355e-28+2.220446049250313e-14j))


The handling of uncertainties is incorrect.

SciPy’s quad_vec() function#

The easiest solution, however, seems to be scipy.integrate.quad_vec():

quad_vec(integrand, 0.0, 2.0)
((2.666666666666667+1.9999999999999996j), 1.1302447741611209e-13)

This has the added benefit that it can handle functions that return arrays:

def gaussian(x, mu, sigma):
    return np.exp(-((x - mu) ** 2) / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi))

mu_values = np.linspace(-2, +3, num=10)
result, _ = quad_vec(lambda x: gaussian(x, mu_values, sigma=0.5), 0, 2.0)
array([3.16712418e-05, 1.93302827e-03, 3.77201760e-02, 2.52491007e-01,
       6.71450766e-01, 9.32839322e-01, 9.04958400e-01, 5.87850435e-01,
       1.87030892e-01, 2.27501310e-02])

Integrate with quadpy#


quadpy now requires a license. The examples below are only shown for documentation purposes.

Alternatively, one could use quadpy, which essentially does the same as in quad(), but can also (to a large degree) handle vectorized input and properly handles uncertainties. For example:

from functools import partial

def parametrized_func(s_prime, s):
    return s_prime * (s_prime + s + 1j)

s_array = np.linspace(-1, 1, num=10)
    partial(parametrized_func, s=s_array),


One may need to play around with the tolerance if the function is not smooth, see sigma-py/quadpy#255.


quadpy raises exceptions with ModuleNotFoundErrors that are a bit unreadable. They are caused by the fact that orthopy and ndim need to be installed separately.

SymPy integral#

The dispersion integral from Eq. (2) in TR-003 features a variable \(s\) that is an argument to the function \(\Sigma_a\). This becomes a challenge when \(s\) gets vectorized (in this case: gets an event-wise numpy.array of invariant masses). It seems that quad_vec() can handle this well though.

def parametrized_func(s_prime, s):
    return s_prime * (s_prime + s + 1j)

s_array = np.linspace(-1, +1, num=10)
    lambda x: parametrized_func(x, s=s_array),
(array([0.66666667+2.j, 1.11111111+2.j, 1.55555556+2.j, 2.        +2.j,
        2.44444444+2.j, 2.88888889+2.j, 3.33333333+2.j, 3.77777778+2.j,
        4.22222222+2.j, 4.66666667+2.j]),

We now attempt to design SymPy expression classes that correctly lambdify() using this vectorized numerical integral for handles complex values. Note that this integral expression class derives from sympy.Integral and that:

  1. overwrites its doit() method, so that the integral cannot be evaluated by SymPy,

  2. provides a custom NumPy printer method (see TR-001) that lambdifies this expression node to quadpy.quad(),

  3. adds class variables that allow configuring quad_vec(),

  4. dummifies the integration variable in case it is not a valid Python variable name.

class UnevaluatableIntegral(sp.Integral):
    abs_tolerance = 1e-5
    rel_tolerance = 1e-5
    limit = 50
    dummify = True

    def doit(self, **hints):
        args = [arg.doit(**hints) for arg in self.args]
        return self.func(*args)

    def _numpycode(self, printer, *args):
        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]
        if self.dummify:
            dummy = sp.Dummy()
            expr = expr.xreplace({x: dummy})
            x = dummy
        integrate_func = "quad_vec"
        return (
            f"{integrate_func}(lambda {printer._print(x)}: {printer._print(expr)},"
            f" {printer._print(a)}, {printer._print(b)},"
            f" epsabs={self.abs_tolerance}, epsrel={self.abs_tolerance},"
            f" limit={self.limit})[0]"

To test whether this works, test this integral expression on another unevaluated() expression:

Hide code cell source
class MyFunction(sp.Expr):
    x: sp.Symbol
    omega1: sp.Symbol
    omega2: sp.Symbol
    phi1: sp.Symbol
    phi2: sp.Symbol
    _latex_repr_ = R"f\left({x}\right)"

    def evaluate(self) -> sp.Expr:
        x, omega1, omega2, phi1, phi2 = self.args
        return sp.sin(omega1 * x + phi1) + sp.sin(omega2 * x + phi2)

x, omega1, omega2, phi1, phi2 = sp.symbols("x omega1 omega2 phi1 phi2")
expr = MyFunction(x, omega1, omega2, phi1, phi2)
Math(aslatex({expr: expr.doit(deep=False)}))
\[\begin{split}\displaystyle \begin{array}{rcl} f\left(x\right) &=& \sin{\left(\omega_{1} x + \phi_{1} \right)} + \sin{\left(\omega_{2} x + \phi_{2} \right)} \\ \end{array}\end{split}\]
w, a, b = sp.symbols("w a b")
fourier_expr = UnevaluatableIntegral(expr * sp.exp(-sp.I * w * x), (x, a, b))
\[\displaystyle \int\limits_{a}^{b} e^{- i w x} f\left(x\right)\, dx\]

Indeed the expression correctly lambdifies correctly, despite the doit() call:

func = sp.lambdify([x, omega1, omega2, phi1, phi2], expr.doit())
fourier_func = sp.lambdify([w, omega1, omega2, phi1, phi2, a, b], fourier_expr.doit())
Hide code cell source
src = inspect.getsource(fourier_func)
src = f"""```python
{black.format_str(src, mode=black.FileMode()).strip()}
def _lambdifygenerated(w, omega1, omega2, phi1, phi2, a, b):
    return quad_vec(
        lambda _Dummy_36: (
            sin(_Dummy_36 * omega1 + phi1) + sin(_Dummy_36 * omega2 + phi2)
        * exp(-_Dummy_36 * 1j * w),
domain = np.linspace(-7, +7, num=500)
parameters = dict(
func_output = func(domain, **parameters)
fourier_output = fourier_func(domain, **parameters, a=-10, b=+10)
Hide code cell source
%config InlineBackend.figure_formats = ['svg']

fig, ax = plt.subplots()
ax.plot(domain, func_output, label="$f(x)$")
ax.plot(domain, fourier_output.real, label=R"$\mathrm{Re}\,F(w)$")
ax.plot(domain, fourier_output.imag, label=R"$\mathrm{Im}\,F(w)$")


See how this integral expression class is applied to the phase space factor in TR-003.