TR-016

Symbolic integral

This report investigates how to formulate a symbolic integral that correctly evaluates to

To be implemented

Lambdifying a symbolic integral

Import Python libraries

Hide code cell source

import inspect

import black
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from ampform.io 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[2], line 5
      1 def integrand(x):
      2     return x * (x + 1j)
----> 5 quad(integrand, 0.0, 2.0)

File ~/work/report/report/docs/016/.venv/lib/python3.12/site-packages/scipy/integrate/_quadpack_py.py:459, in quad(func, a, b, args, full_output, epsabs, epsrel, limit, points, weight, wvar, wopts, maxp1, limlst, complex_func)
    456     return retval
    458 if weight is None:
--> 459     retval = _quad(func, a, b, args, full_output, epsabs, epsrel, limit,
    460                    points)
    461 else:
    462     if points is not None:

File ~/work/report/report/docs/016/.venv/lib/python3.12/site-packages/scipy/integrate/_quadpack_py.py:606, in _quad(func, a, b, args, full_output, epsabs, epsrel, limit, points)
    604 if points is None:
    605     if infbounds == 0:
--> 606         return _quadpack._qagse(func,a,b,args,full_output,epsabs,epsrel,limit)
    607     else:
    608         return _quadpack._qagie(func, bound, infbounds, args, full_output, 
    609                                 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))

Warning

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)
result

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

Warning

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)
quadpy.quad(
    partial(parametrized_func, s=s_array),
    a=0.0,
    b=2.0,
)

Note

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

Tip

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)
quad_vec(
    lambda x: parametrized_func(x, s=s_array),
    a=0.0,
    b=2.0,
)

(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]),
 3.8263293934812236e-13)

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"
        printer.module_imports["scipy.integrate"].add(integrate_func)
        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:

Show code cell source

Hide code cell source

@unevaluated
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))
fourier_expr

\[\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())

Show code cell source

Hide code cell source

src = inspect.getsource(fourier_func)
src = f"""```python
{black.format_str(src, mode=black.FileMode()).strip()}
```"""
Markdown(src)

def _lambdifygenerated(w, omega1, omega2, phi1, phi2, a, b):
    return quad_vec(
        lambda _Dummy_34: (
            sin(_Dummy_34 * omega1 + phi1) + sin(_Dummy_34 * omega2 + phi2)
        )
        * exp(-_Dummy_34 * 1j * w),
        a,
        b,
        epsabs=1e-05,
        epsrel=1e-05,
        limit=50,
    )[0]

domain = np.linspace(-7, +7, num=500)
parameters = dict(
    omega1=1.2,
    omega2=2.3,
    phi1=-1.2,
    phi2=+0.4,
)
func_output = func(domain, **parameters)
fourier_output = fourier_func(domain, **parameters, a=-10, b=+10)

Show code cell source

Hide code cell source

%config InlineBackend.figure_formats = ['svg']

fig, ax = plt.subplots()
ax.set_xlabel("$x,w$")
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)$")
ax.legend()
plt.show()

Tip

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