Search code examples
numpycythonpython-c-api

passing 1 or 2 d numpy array to c throw cython


I am writing an extension to my python code in c and cython, by following this guide.

my c function signature is

void c_disloc(double *pEOutput, double *pNOutput, double *pZOutput, double *pModel, double *pECoords, double *pNCoords, double nu, int NumStat, int NumDisl)

and my cython function is

cdef extern void c_disloc(double *pEOutput, double *pNOutput, double *pZOutput, double *pModel, double *pECoords, double *pNCoords, double nu, int NumStat, int NumDisl)

@cython.boundscheck(False)
@cython.wraparound(False)
def disloc(np.ndarray[double, ndim=2, mode="c"] pEOutput not None,
           np.ndarray[double, ndim=2, mode="c"] pNOutput not None,
           np.ndarray[double, ndim=2, mode="c"] pZOutput not None,
           np.ndarray[double, ndim=1, mode="c"] pModel not None,
           np.ndarray[double, ndim=2, mode="c"] pECoords not None,
           np.ndarray[double, ndim=2, mode="c"] pNCoords not None,
           double nu,int NumStat, int NumDisl ):

    c_disloc(&pEOutput[0,0], &pNOutput[0,0], &pZOutput[0,0], &pModel[0], &pECoords[0,0], &pNCoords[0,0], nu, NumStat, NumDisl)

    return None

now my c function has the same behavior no matter if the arrays that its getting are 1d or 2d arrays, but I didn't succeed making the cython function to be able to get 1d or 2d numpy arrays. of course, I could write tow cython function one for the 1d case and one for the 2d case but it will be cleaner to do it with one function. dose someone knows how to do it?


Solution

  • I'd accept an untyped argument, check that it's a C contiguous array and then use np.ravel to get a flat array (this returns a view, not a copy, when passed a C contiguous array). It's easy to create that as a cdef function:

    cdef double* get_array_pointer(arr) except NULL:
        assert(arr.flags.c_contiguous) # if this isn't true, ravel will make a copy
        cdef double[::1] mview = arr.ravel()
        return &mview[0]
    

    Then you'd do

    def disloc(pEOutput,
               pNOutput,
               # etc...
               double nu,int NumStat, int NumDisl ):
    
        c_disloc(get_array_pointer(pEOutput), get_array_pointer(pNOutput),
                # etc
                nu, NumStat, NumDisl)
    

    I've removed the

    @cython.boundscheck(False)
    @cython.wraparound(False)
    

    since it's obvious they will gain you close to nothing. Using them without thinking about whether they do anything seems like cargo cult programming to me.