Lambdifying a symbolic integral#
Import Python libraries
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:
overwrites its
doit()method, so that the integral cannot be evaluated by SymPy,provides a custom NumPy printer method (see TR-001) that lambdifies this expression node to
quadpy.quad(),adds class variables that allow configuring
quad_vec(),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
@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)}))
w, a, b = sp.symbols("w a b")
fourier_expr = UnevaluatableIntegral(expr * sp.exp(-sp.I * w * x), (x, a, b))
fourier_expr
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
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
%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.