如何将附加参数传递给 numba cfunc 作为 LowLevelCallable 传递给 scipy.integrate.quad

Question

文档 discusses 使用 numba 的 cfunc 作为 scipy.integrate.quad 的 LowLevelCallable 参数。我需要同样的东西和额外的参数。

我基本上是在尝试做这样的事情：

import numpy as np
from numba import cfunc
import numba.types
voidp = numba.types.voidptr
def integrand(t, params):
    a = params[0] # this is additional parameter
    return np.exp(-t/a) / t**2
nb_integrand = cfunc(numba.float32(numba.float32, voidp))(integrand)

然而，它不起作用，因为params应该是voidptr/void*，它们不能转换为double。我收到以下错误消息：

TypingError: Failed at nopython (nopython frontend)
Invalid usage of getitem with parameters (void*, int64)
 * parameterized

我没有找到任何关于如何在 Numba 中从 void* 中提取值的信息。在 C 中，它应该类似于 a = *((double*) params) — 是否可以在 Numba 中做同样的事情？

Answer 1

1.通过 scipy.integrate.quad

传递额外参数

quad docs 说：

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

double func(double x)

double func(double x, void *user_data)

double func(int n, double *xx)

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

The user_data is the data contained in the scipy.LowLevelCallable. In the call forms with xx, n is the length of the xx array which contains xx[0] == x and the rest of the items are numbers contained in the args argument of quad.

因此，要通过 quad 向 integrand 传递一个额外的参数，您最好使用 double func(int n, double *xx) 签名。

您可以为被积函数编写装饰器，将其转换为 LowLevelCallable，如下所示：

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 integrand(t, *args):
    a = args[0]
    return np.exp(-t/a) / t**2

def do_integrate(func, a):
    """
    Integrate the given function from 1.0 to +inf with additional argument a.
    """
    return si.quad(func, 1, np.inf, args=(a,))

print(do_integrate(integrand, 2.))
>>>(0.326643862324553, 1.936891932288535e-10)

或者如果您不需要装饰器，请手动创建 LowLevelCallable 并将其传递给 quad。

2。包装被积函数

我不确定以下内容是否满足您的要求，但您也可以包装 integrand 函数以获得相同的结果：

import numpy as np
from numba import cfunc
import numba.types

def get_integrand(*args):
    a = args[0]
    def integrand(t):
        return np.exp(-t/a) / t**2
    return integrand

nb_integrand = cfunc(numba.float64(numba.float64))(get_integrand(2.))

import scipy.integrate as si
def do_integrate(func):
    """
    Integrate the given function from 1.0 to +inf.
    """
    return si.quad(func, 1, np.inf)

print(do_integrate(get_integrand(2)))
>>>(0.326643862324553, 1.936891932288535e-10)
print(do_integrate(nb_integrand.ctypes))
>>>(0.326643862324553, 1.936891932288535e-10)

3。从 voidptr 转换为 python 类型

我认为这还不可能。从 2016 年的 this discussion 开始，voidptr 似乎只是为了将上下文传递给 C 回调。

The void * pointer case would be for APIs where foreign C code does not every try to dereference the pointer, but simply passes it back to the callback as way for the callback to retain state between calls. I don't think it is particularly important at the moment, but I wanted to raise the issue.

并尝试以下操作：

numba.types.RawPointer('p').can_convert_to(
    numba.typing.context.Context(), CPointer(numba.types.Any)))
>>>None

似乎也不令人鼓舞！

Answer 2

此处采用与 Jacques Gaudin 建议的第一点相同的技术，但有几个论点。

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):
        values = carray(xx, n)
        return jitted_function(values)
    return LowLevelCallable(wrapped.ctypes)

@jit_integrand_function
def integrand(args):
    t = args[0]
    a = args[1]
    b = args[2]
    return b * np.exp(-t/a) / t**2

def do_integrate(func, a):
    """
    Integrate the given function from 1.0 to +inf with additional argument a.
    """
    return si.quad(func, 1, np.inf, args=(a, b,))

如何将附加参数传递给 numba cfunc 作为 LowLevelCallable 传递给 scipy.integrate.quad

How to pass additional parameters to numba cfunc passed as LowLevelCallable to scipy.integrate.quad

python

scipy

numba