如何从Numba实现C callable与nquad的高效集成？

Question

我需要在python中进行6D的数值积分.因为scipy.integrate.nquad函数很慢,所以我现在正试图通过将整数定义为带有Numba的scipy.LowLevelCallable来加快速度.

我能够通过复制这里给出的示例在1D中使用scipy.integrate.quad执行此操作:

import numpy as np
from numba import cfunc
from scipy import integrate

def integrand(t):
    return np.exp(-t) / t**2

nb_integrand = cfunc("float64(float64)")(integrand)

# regular integration
%timeit integrate.quad(integrand, 1, np.inf)

10000循环,最佳3:每循环128μs

# integration with compiled function
%timeit integrate.quad(nb_integrand.ctypes, 1, np.inf)

100000个循环,最佳3:每循环7.08μs

当我想用nquad做这个时,nquad文档说:

如果用户希望提高集成性能,则f可能是scipy.LowLevelCallable,其中一个签名:

double func(int n, double *xx)
double func(int n, double *xx, void *user_data)

其中n是额外参数的数量,args是附加参数的双精度数组,xx数组包含坐标.user_data是scipy.LowLevelCallable中包含的数据.

但是下面的代码给出了一个错误:

from numba import cfunc
import ctypes

def func(n_arg,x):
    xe = x[0]
    xh = x[1]
    return np.sin(2*np.pi*xe)*np.sin(2*np.pi*xh)

nb_func = cfunc("float64(int64,CPointer(float64))")(func)

integrate.nquad(nb_func.ctypes, [[0,1],[0,1]], full_output=True)

error:quad:第一个参数是一个带有错误签名的ctypes函数指针

是否可以使用numba编译一个函数,该函数可以直接在代码中与nquad一起使用,而无需在外部文件中定义函数？

非常感谢你提前!

Answer 1

将函数包装在 a 中scipy.LowLevelCallable会让人nquad高兴：

si.nquad(sp.LowLevelCallable(nb_func.ctypes), [[0,1],[0,1]], full_output=True)
# (-2.3958561404687756e-19, 7.002641250699693e-15, {'neval': 1323})