How to pass additional parameters to numba cfunc passed as LowLevelCallable to scipy.integrate.quad
Solution 1:
1. Passing extra arguments through scipy.integrate.quad
The quad
docs say:
If the user desires improved integration performance, then
f
may be ascipy.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 thescipy.LowLevelCallable
. In the call forms withxx
,n
is the length of thexx
array which containsxx[0] == x
and the rest of the items are numbers contained in theargs
argument ofquad
.
Therefore to pass an extra argument to integrand
through quad
, you are better of using the double func(int n, double *xx)
signature.
You can write a decorator to your integrand function to transform it to a LowLevelCallable
like so:
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)
Or if you don't want the decorator, create the LowLevelCallable
manually and pass it to quad
.
2. Wrapping the integrand function
I am not sure if the following would meet your requirements but you could also wrap your integrand
function to achieve the same result:
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. Casting from voidptr
to a python type
I don't think this is possible yet. From this discussion in 2016, it seems that voidptr
is only here to pass a context to a C callback.
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.
And trying the following:
numba.types.RawPointer('p').can_convert_to(
numba.typing.context.Context(), CPointer(numba.types.Any)))
>>>None
doesn't seem encouraging either!
Solution 2:
Here the same technique as the first point suggested by Jacques Gaudin, but for several arguments.
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,))