Cython ndarray с произвольными размерами

Question

Я пишу код Cython, который должен обрабатывать NumPy ndarrays с произвольным числом измерений. В настоящее время я просто различные функции, которые принимают ndarrays разных размеров, вроде:

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. 

Но я хотел бы написать один обобщенную функцию, которая принимает ndarray произвольной размерности в качестве аргумента. Я попробовал просто оставить аргумент «ndim», но тогда Cython предполагает ndim = 1, и это нехорошо.

Есть ли способ сделать это, или мне просто нужно написать одну функцию для каждого числа измерений?

Answer 1

Если вы просто хотите сделать что-то поэлементно, хитрость заключается в том, чтобы просто получить 1D вид на массив и работать на том, что

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