How to combine scipy interp1d with mpmath quadosc

I have a density function (from quantum mechanics calculations) to be multiplied with the spherical Bessel function with a momentum grid (momentum q 1d array, real space distance r 1d array, so need to calculate jn(q*r) 2d array). The product will be integrated across the real space to get results as function of momentum (results in 1d array same shape as q).

The Bessel function has oscillation, while the density function fast decay over a threshold distance. I used the adaptive quadrature in quadpy which is fine when oscillation is slow but it fails with high oscillation (for high momentum values or high orders in Bessel functions). The mpmath quadosc could be a nice option, but currently I have the problem that "object arrays are not supported", which seems to be the same as in Relation between mpmath and scipy: Type Error, what would be the best way to solve it since the density function is calculated outside of the mpmath.

from mpmath import besselj, sqrt, pi, besseljzero, inf,quadosc
from scipy.interpolate import interp1d

n = 1
q = np.geomspace(1e-7, 500, 1000)
# lets create a fake gaussian density 
x = np.geomspace(1e-7, 10, 1000)
y = np.exp(-(x-5)**2)
density = interp1d(x,y,kind='cubic',fill_value=0,bounds_error=False)


# if we just want to integrate the spherical bessel function
def spherical_jn(x,n=n):
    return besselj(n + 1 / 2, x) * sqrt(pi / 2 / x)
# this is fine
vals = quadosc(
    spherical_jn, [0, inf], zeros=lambda m: besseljzero(n + 1 / 2, m)
)

# now we want to integrate the spherical bessel function times the density
def spherical_jn_density(x,n=lprimeprime):
    grid = q[..., None] *x
    return besselj(n + 1 / 2,  grid) * sqrt(pi / 2 / grid)*density(x)

# this will fail
vals_density = quadosc(
    spherical_jn_density, [0, inf], zeros=lambda m: besseljzero(n + 1 / 2, m)
)

Expect: an accurate integral of highly oscillating spherical Bessel function with arbitrary decaying function (decay to zero at large distance).

Solution

Your density is interp callable, which works like:

In [33]: density(.5)
Out[33]: array(1.60522789e-09)

It does not work when given a mpmath object:

In [34]: density(mpmath.mpf(.5))
ValueError: object arrays are not supported

It's ok if x is first converted to ordinary float:

In [37]: density(float(mpmath.mpf(.5)))
Out[37]: array(1.60522789e-09)

Tweaking your function:

def spherical_jn_density(x,n=1):
    print(repr(x))
    grid = q[..., None] *x
    return besselj(n + 1 / 2,  grid) * sqrt(pi / 2 /grid) * density(x)

and trying to run the quadosc (with a smaller q)

In [57]: vals_density = quadosc(
    ...:     spherical_jn_density, [0, inf], zeros=lambda m: besseljzero(n + 1 / 2, m))
mpf('0.506414729137261838698106')
TypeError: cannot create mpf from array([[mpf('5.06414729137261815781894e-8')],
       [mpf('0.000000473559111442409924364745')],
       [mpf('0.00000442835129247081824275722')],
       [mpf('0.0000414104484439061558283487')],
       [mpf('0.000387237851532012775822723')],
       [mpf('0.00362114295531604773233197')],
       [mpf('0.0338620727569835882851491')],
       [mpf('0.316651395857188250996884')],
       [mpf('2.96107409661232278850947')],
       [mpf('27.6896294168963266721213')]], dtype=object)

In other words,

besselj(n + 1 / 2,  grid)

is having problems, even before trying to evaluate density(x). mpmath functions don't work with numpy arrays; and many numpy/scipy functions don't work with mpmath objects.