如何使用Numba在SciPy中使用任意数量的变量和参数执行多重集成？

Question

我想使用 Numba 装饰多重积分的被积，以便 SciPy的Nquad 函数可以将其作为 LowLevelCallable 调用

强>。理想情况下，装饰器应该允许Nquad的 args 参数中包含任意数量的变量和任意数量的其他参数。这是基于出色的Q&A from earlier this year构建的，但是扩展到了多个变量和参数的情况。

作为一个例子，假设以下具有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)

理想的代码在 integrand 函数上仅使用一个装饰器即可运行：

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

Numba documents引用一个carray函数，当在回调中给出低级指针和数组大小时，该函数应返回Numpy数组。可能的是，这可以用于推广超出二变量二参数情况的代码。我（未成功）尝试实现这一点的是在两行注释掉的代码行中。

我们将不胜感激。的确，一位Numba开发人员pointed out认为SciPy集成是Numba编写的原因之一，但缺少该领域的文档和示例。

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)