Calling Numba cfunc from njitted function with numpy array argument

Question

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)

Nin17 · Accepted Answer · 2024-11-01 21:09:49Z

1

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)

edited Nov 1, 2024 at 21:09

answered Nov 1, 2024 at 20:53

Nin17

3,6022 gold badges7 silver badges18 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

kesh Over a year ago

Thanks, this worked! It's unfortunate that you cannot do the same for non-jitted function (only for the benchmarking purpose, so it's OK).

Nin17 Over a year ago

well if you really wanted to have the same code, you could overload the .data_as method

kesh Over a year ago

True. I'll give it a thought. Thx

Collectives™ on Stack Overflow

Calling Numba cfunc from njitted function with numpy array argument

1 Answer 1

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related