python - 您如何实现可从 Numba 调用的 C 以与 nquad 高效集成？

Question

我需要在 python 中进行 6D 数值积分。因为 scipy.integrate.nquad 函数很慢，我目前正在尝试通过将被积函数定义为带有 Numba 的 scipy.LowLevelCallable 来加快速度。

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

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 微秒

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

100000 次循环，最好的 3 次:每次循环 7.08 微秒

当我现在想用 nquad 执行此操作时，nquad 文档说:

If the user desires improved integration performance, then f may be a scipy.LowLevelCallable with one of the signatures:

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

where n is the number of extra parameters and args is an array of doubles of the additional parameters, the xx array contains the coordinates. The user_data is the data contained in the 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)

错误:quad:第一个参数是签名不正确的 ctypes 函数指针

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

非常感谢您!

Answer 1

将函数包装在 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})