如何使用Numba在SciPy中执行多个积分，并使用任意数量的变量和参数？

Question

我想使用Numba来修饰多重积分的被积函数，以便可以通过SciPy的Nquad函数调用它作为LowLevelCallable。理想情况下，装饰器应该允许任意数量的变量和来自Nquad的args参数的任意数量的附加参数。这是基于今年早些时候的一个出色的问答构建的，但扩展到了多个变量和参数的情况。

例如，假设以下具有N个变量和K个参数的多重积分：

以下代码可以工作，但仅适用于两个变量和两个参数（N=2，K=2）。它不适用于更一般的情况。这是因为装饰器中的某些参数是手动枚举的（在wrapped函数内部的xx[0]、xx[1]、xx[2]、xx[3]）。装饰器必须为每个不同数量的变量或参数进行编辑。如果可能的话，我想避免这种情况。请注意，integrand函数本身利用了Numpy对象和方法，因此没有这个问题。

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)

理想的代码只需要在“被积函数”上使用一个装饰器就可以运行：

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

Numba 文档 中提到了一个 carray 函数，该函数应在回调中给出数组的低级指针和大小时返回一个Numpy数组。可能可以使用它来将代码推广到两个变量两个参数之外的情况。我（不成功地）尝试实现这个功能在两行被注释掉的代码中。

需要帮助。事实上，Numba开发人员指出，SciPy集成是编写Numba的原因之一，但在这个领域缺乏文档和示例。

- Aboottogo

Answer 1

以下代码可行：

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)

- Aboottogo