Как передать дополнительные параметры numba cfunc, переданные как LowLevelCallable в scipy.integrate.quad

Документация обсуждает с использованием Numba cfunc в качестве LowLevelCallable аргумента scipy.integrate.quad. Мне нужно то же самое с дополнительным параметром.

Я в основном пытаюсь сделать что-то вроде этого:

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

Я не нашел никакой информации о том, как извлекать значения из void* в Numba. В C это должно быть что-то вроде a = *((double*) params) - возможно ли сделать то же самое в Numba?

Ответ 1

1. Передача дополнительных аргументов через scipy.integrate.quad

quad документы говорят:

Если пользователь желает улучшить производительность интеграции, то f может быть scipy.LowLevelCallable с одной из подписей:

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)

user_data - это данные, содержащиеся в scipy.LowLevelCallable. В формах вызовов с xx n - это длина массива xx которая содержит xx[0] == x а остальные элементы - числа, содержащиеся в аргументе args quad.

Поэтому, чтобы передать дополнительный аргумент в integrand через quad, вам лучше использовать 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 году кажется, что voidptr только здесь, чтобы передать контекст для обратного вызова C.

Случай с указателем void * был бы для API, где внешний код C не каждый пытается разыменовать указатель, но просто передает его обратно на обратный вызов, как способ для обратного вызова сохранить состояние между вызовами. Я не думаю, что это особенно важно на данный момент, но я хотел поднять вопрос.

И попробуйте следующее:

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

тоже не радует!