Cython ndarray with arbitrary dimensions

Question

I'm writing some Cython code that needs to be able to handle NumPy ndarrays that have an arbitrary number of dimensions. Currently, I just have different functions that accept ndarrays of different sizes, sort of like:

def func1(np.ndarray[DTYPE_float64_t, ndim=1] arr):
    # Do something with the 1-D ndarray.

def func2(np.ndarray[DTYPE_float64_t, ndim=2] arr):
    # Do something with the 2-D ndarray.

def func3(np.ndarray[DTYPE_float64_t, ndim=3] arr):
    # Do something with the 3-D ndarray.

But I would like to write a single generic function that takes an ndarray of arbitrary dimension as an argument. I tried simply leaving the "ndim" argument off, but then Cython assumes ndim=1, and that's no good.

Is there a way to do this, or do I just have to write one function for each number of dimensions?

Answer 1

If you're just looking to do something elementwise, the trick is to just do get a 1D view of the array and operate on that

def func(arr):
   shape = arr.shape
   output = _func_impl(arr.ravel())
   return output.reshape(shape) # ensure that the output is the same shape
       # as the input. Skip this if it doesn't make sense!

def _func_impl(np.ndarray[DTYPE_float64_t, ndim=1] arr):
   # do something useful