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

Question

我想用Numba来装饰一个多重积分的积分,这样它就可以被SciPy的Nquad函数称为LowLevelCallable.理想情况下,装饰器应该允许任意数量的变量,以及来自Nquad的args参数的任意数量的附加参数.这是从今年早些时候开始的优秀问答,但扩展到多个变量和参数的情况.

例如,假设以下与N个变量和K参数的多个积分:

以下代码有效,但仅适用于两个变量和两个参数(N = 2,K = 2).它不适用于更一般的情况.这是因为装饰器中的一些参数是手动枚举的(包裹函数内的xx [0],xx [1],xx [2],xx [3] ).必须针对每个不同数量的变量或参数编辑装饰器.如果可能的话,我想避免这种情况.请注意,被积函数本身利用了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 文档引用了一个carray函数,它应该在给出回调中的低级指针和数组大小时返回Numpy数组.可能,这可以用于概括超出双变量二参数情况的代码.我(不成功)尝试实现这一点是在两个注释掉的代码行中.

帮助将不胜感激.事实上,Numba的一位开发人员指出,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)