How to use Numba to perform multiple integration in SciPy with an arbitrary number of variables and parameters?

Question

I'd like to use Numba to decorate the integrand of a multiple integral so that it can be called by SciPy's Nquad function as a LowLevelCallable. Ideally, the decorator should allow for an arbitrary number of variables, and an arbitrary number of additional parameters from the Nquad's args argument. This is built off an excellent Q&A from earlier this year, but extended to the case of multiple variables and parameters.

As an example, suppose the following multiple integral with N variables and K parameters:

The following code works, but only for two variables and two parameters (N=2,K=2). It does not work for the more general case. This is because some of the arguments in the decorator are manually enumerated (xx[0],xx[1],xx[2],xx[3] inside the wrapped function). The decorator would have to be edited for every different number of variables or parameters. I'd like to avoid that, if possible. Note that the integrand function itself takes advantage of Numpy objects and methods and so does not have this problem.

import numpy as np
import scipy.integrate as si
import numba
from numba import cfunc,carray
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], xx[2], xx[3])
        #xx = carray(xx,len(xx))
        #return jitted_function(xx)
    return LowLevelCallable(wrapped.ctypes)

@jit_integrand_function
def integrand(*args):
    d = np.array([args])
    return -np.exp(d.prod())

#Two variable, two parameter example
parms = np.array([2,3])
print si.nquad(integrand,[[0,1],[0,1]],parms)

The ideal code would use just one decorator on the integrand function to also run:

#Three variable, three parameter example
parms2 = np.array([1,2,3])
print si.nquad(integrand,[[0,1],[0,1],[0,1]],parms2)

The Numba documents refer to a carray function that ought to return a Numpy array when given the low-level pointer and size of an array in the callback. Possibly, this could be used to generalize the code beyond the two-variable-two-parameter case. My (unsuccessful) attempt to implement this is in the two commented-out lines of code.

Help would be appreciated. Indeed, one of the Numba developers pointed out that SciPy integration was one of the reasons Numba was written, but that documentation and examples in this area are lacking.

Answer 1

The following code works:

import numpy as np
import scipy.integrate as si
import numba
from numba import cfunc,carray
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):
        values = carray(xx,n)
        return jitted_function(values)
    return LowLevelCallable(wrapped.ctypes)

@jit_integrand_function
def integrand(args):
    return -np.exp(args.prod())

#Two variable, two parameter example
parms = np.array([2,3])
print si.nquad(integrand,[[0,1],[0,1]],parms)

#Three variable, three parameter example
parms2 = np.array([1,2,3])
print si.nquad(integrand,[[0,1],[0,1],[0,1]],parms2)