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

Question

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

通过复制给定的示例 here:

，我能够使用 scipy.integrate.quad 在 1D 中执行此操作

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})

Answer 2

您传递给 nquad 的函数的签名应该是 double func(int n, double *xx)。您可以像这样为函数 func 创建装饰器：

import numpy as np
import scipy.integrate as si
import numba
from numba import cfunc
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])
    return LowLevelCallable(wrapped.ctypes)

@jit_integrand_function
def func(xe, xh):
    return np.sin(2*np.pi*xe)*np.sin(2*np.pi*xh)

print(si.nquad(func, [[0,1],[0,1]], full_output=True))
>>>(-2.3958561404687756e-19, 7.002641250699693e-15, {'neval': 1323})

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

How can you implement a C callable from Numba for efficient integration with nquad?

python

numpy

scipy

numba