CODEWITHSUNDEEP
pythonnumba
kesh
kesh

Reputation: 5463

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)

Upvotes: 0

Views: 48

Answers (1)

Nin17
Nin17

Reputation: 3472

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)

Upvotes: 1

Related Questions