pythonnumba

Calling Numba cfunc from njitted function with numpy array argument


I'm trying to call a cfunction inside njitted function, but Numba does not have data_as() method for its array to cast double pointer. Could anyone help me figure out how to make it work?

import ctypes
import numpy as np
import numba as nb

@nb.cfunc(nb.types.void(
                nb.types.CPointer(nb.types.double),
                nb.types.CPointer(nb.types.double),
                nb.types.int64,
                nb.types.int64,
                nb.types.int64,
                nb.types.int64,
            ))
def get_param2(xn_, x_, idx, n, m1, m2):
    in_array = nb.carray(x_, (n, m1, m2))
    out_array = nb.carray(xn_, (m1, m2))
    if idx >= n:
        idx = n - 1
    out_array[:, :] = in_array[idx]


def test_get_param(): # this one works

    A = np.zeros((100, 2, 3))
    Ai = np.empty((2, 3))
    get_param2(
        Ai.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        A.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        40,
        *A.shape,
    )

    assert np.array_equal(A[40], Ai)


@nb.jit(nopython=True)
def get_param_njit(A, i):
# this one fails with the error:
#     `Unknown attribute 'data_as' of type ArrayCTypes(dtype=float64, ndim=2)`

    Ai = np.empty((2, 3))
    get_param2(
        Ai.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        A.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), 
        i, 
       *A.shape)

    return Ai


def test_get_param_njit():
    A = np.zeros((100, 2, 3))
    Ai = get_param_njit(A, 40)
    assert np.array_equal(A[40], Ai)

Solution

  • You just use the .ctypes attribute of the array:

    @nb.njit
    def test_get_param_njit(): 
    
        A = np.zeros((100, 2, 3))
        Ai = np.empty((2, 3))
        get_param2(
            Ai.ctypes,
            A.ctypes,
            np.intc(40), # Convert to intc for your signature
            np.intc(100), np.intc(2), np.intc(3),
        )
    
        assert np.array_equal(A[40], Ai)