如何使用 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)

理想的代码会在 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 refer to a carray 函数应该 return 一个 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)

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

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

python

numpy

scipy

numba