Is there a way to use scipy.optimize.fsolve with jit_integrand_function and scipy.integrate.quad?

Question

Based on the explanation provided here 1, I am trying to use the same idea to speed up the following integral:

import scipy.integrate as si
from scipy.optimize import root, fsolve
import numba
from numba import cfunc
from numba.types import intc, CPointer, float64
from scipy import LowLevelCallable

def integrand(t, *args):
    a = args[0]
    c = fsolve(lambda x: a * x**2 - np.exp(-x**2 / a), 1)[0]
    return c * np.exp(- (t / (a * c))**2) 

def do_integrate(func, a):
    return si.quad(func, 0, 1, args=(a,))

print(do_integrate(integrand, 2.)[0])

With the previous reference, I tried to use numba/jit and modify the previous block in the following way:

import numpy as np
import scipy.integrate as si
from scipy.optimize import root
import numba
from numba import cfunc
from numba.types import intc, CPointer, float64
from scipy import LowLevelCallable

def jit_integrand_function(integrand_function):
    jitted_function = numba.jit(integrand_function, nopython=True)  
    @cfunc(float64(intc, CPointer(float64)))
    def wrapped(n, xx):
        return jitted_function(xx[0], xx[1])
    return LowLevelCallable(wrapped.ctypes)

@jit_integrand_function
def integrand(t, *args):
    a = args[0]
    c = fsolve(lambda x: a * x**2 - np.exp(-x**2 / a), 1)[0]
    return c * np.exp(- (t / (a * c))**2)

def do_integrate(func, a):
    return si.quad(func, 0, 1, args=(a,))

do_integrate(integrand, 2.)

However, this implementation gives me the error


TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: convert make_function into JIT functions)
Cannot capture the non-constant value associated with variable 'a' in a function that will escape.

File "<ipython-input-16-3d98286a4be7>", line 20:
def integrand(t, *args):
    <source elided>
    a = args[0]
    c = fsolve(lambda x: a * x**2 - np.exp(-x**2 / a), 1)[0]
    ^

During: resolving callee type: type(CPUDispatcher(<function integrand at 0x11a949d08>))
During: typing of call at <ipython-input-16-3d98286a4be7> (14)

During: resolving callee type: type(CPUDispatcher(<function integrand at 0x11a949d08>))
During: typing of call at <ipython-input-16-3d98286a4be7> (14)

The error is coming from the fact that I am using fsolve from scipy.optimize inside the integrand function.

I would like to know if there is a workaround this error and if it is possible to use the scipy.optimize.fsolve with numba in this context.

It is unfortunately not possible to apply Numba's njit decorator on code containing Scipy function calls. — Alex van Houten, Aug 21 '23 at 09:05

score 0 · Answer 1 · answered May 09 '21 at 07:45

I wrote a little python wrapper for Minpack, called NumbaMinpack, which can be called within numba compiled functions: https://github.com/Nicholaswogan/NumbaMinpack. You can use it to @njit the integrand:

import scipy.integrate as si
from NumbaMinpack import hybrd, minpack_sig
from numba import njit, cfunc
import numpy as np

@cfunc(minpack_sig)
def f(x, fvec, args):
    a = args[0]
    fvec[0] = a * x[0]**2.0 - np.exp(-x[0]**2.0 / a)

funcptr = f.address # pointer to function  

@njit
def integrand(t, *args):
    a = args[0]
    args_ = np.array(args)
    x_init = np.array([1.0])
    sol = hybrd(funcptr,x_init,args_)
    c = sol[0][0]
    return c * np.exp(- (t / (a * c))**2) 

def do_integrate(func, a):
    return si.quad(func, 0, 1, args=(a,))

print(do_integrate(integrand, 2.)[0])

On my computer, the above code takes 87 µs, while the pure python version takes 2920 µs

Very interesting. Would that also work with Numba's guvectorize decorator? And with target='cuda'? Can NumbaMinpack do more general things than root finding, e.g. Gauss fitting? — Alex van Houten, Aug 26 '23 at 12:59

Is there a way to use scipy.optimize.fsolve with jit_integrand_function and scipy.integrate.quad?

1 Answers1