How can I return a C struct with numba?

Numba Community Support
1

Hi,

I’m using numba in conjunction with scipy.LowLevelCallable to pass small callables to numerical integrators.

The following MWE works as expected. It allows me to jit a parametrized function and return the result as a scipy.LowLevelCallable


import cffi
import numba
from numba.core.typing import cffi_utils
import scipy

numba_kwargs = {
    'nopython':True     ,
    'cache':True        ,
    'fastmath':True     ,
    'nogil':True        ,
}

src = """
typedef  double (*fun_t)(double) ;
"""

ffi = cffi.FFI()
ffi.cdef(src)

sig = cffi_utils.map_type(ffi.typeof('fun_t'), use_record_dtype=True)

def param_fun(n):

    @numba.cfunc(sig)
    @numba.jit(signature=sig, **numba_kwargs)
    def fun(x):
        return n*x

    return scipy.LowLevelCallable(fun.ctypes)

toto = param_fun(0.5)

Now, how can I do the same, but returning a C struct ?

src = """
typedef struct {
  double a;
} a_t;

typedef a_t (*fun_t)(double) ;
"""

ffi = cffi.FFI()
ffi.cdef(src)

sig = cffi_utils.map_type(ffi.typeof('fun_t'), use_record_dtype=True)

def param_fun(n):

    @numba.cfunc(sig)
    @numba.jit(signature=sig, **numba_kwargs)
    def param_fun(x):
        # ?????

        return n*x

    return scipy.LowLevelCallable(param_fun.ctypes)
2

Hey @gabrielfougeron ,

If you intend to pass only scalars as additional arguments for a SciPy quadrature, you can define your C structure using make_c_struct and utilize the LowLevelCallable function.

import numpy as np
import numba as nb
from numba import types
from scipy import integrate, LowLevelCallable
import ctypes

# Define the C-struct
args_dtype = types.Record.make_c_struct([
    ('x2', types.float64),
])

# function to integrate
def integrand(x1, x2):
    return np.exp(-x1/x2) / x1**2

# function factory
def create_jit_integrand_function(integrand_function, args_dtype):
    jitted_function = nb.njit(integrand_function)

    # double func(double x, void *user_data)
    @nb.cfunc(types.float64(types.float64, types.CPointer(args_dtype)))
    def wrapped(x1, user_data_p):
        user_data = nb.carray(user_data_p, 1)
        x2 = user_data[0].x2
        return jitted_function(x1, x2)
    return wrapped

# wrapped function call
def do_integrate(func, args, a=0, b=1):
    integrand_func = LowLevelCallable(func.ctypes, user_data=args.ctypes.data_as(ctypes.c_void_p))
    # quadrature(func, a, b, args=(), tol=1.49e-8, rtol=1.49e-8, maxiter=50, vec_func=True, miniter=1)
    return integrate.quad(integrand_func, a, b)

# How to use it:
x2 = 1.0
args = np.array((x2,), dtype=args_dtype)
func = create_jit_integrand_function(integrand, args_dtype)
print(do_integrate(func, args, a=1, b=np.inf))
# (0.14849550677592208, 3.8736750296130505e-10)

If you need to define arrays as arguments, there is an example by @max9111

3

Thanks, I did not know about types.Record.make_c_struct