python - 为 n 维数组推广 Cython 代码

Question

我正在编写很多需要泛化到n维的 Cython 代码（大多数情况下这将是n = 1、2 或 3）。

以下面的代码为例，这是对元素的简单求和。

cimport cython
from cython cimport floating

cpdef floating nd_sum(floating [:, :, :] arr):
    cdef:
        Py_ssize_t [3] N = [arr.shape[0], arr.shape[1], arr.shape[2]]
        Py_ssize_t [3] i
        floating total = 0

    for i[0] in range(N[0]):
        for i[1] in range(N[1]):
            for i[2] in range(N[2]):
                total += arr[i[0], i[1], i[2]]

    return total

从概念上讲，对n维的泛化是非常明显的。但是，我似乎无法在代码中实现它。我讨厌必须为n = 1,2,3 的每种情况复制此代码，...

一般来说，我还需要能够访问每个维度中每个数组元素的邻居（想想n维内核卷积等），这使得像展平数组这样的事情变得不可行。

我的主要问题是：

• 如何告诉 Cython 期望具有任意维度的数组作为输入？我想我可能不得不求助于np.ndarray这里......

• 给定n 个索引的数组，如何从n维数组中获取元素？

• 如何概括n嵌套循环的构造？递归在这里似乎是不可避免的，但是我如何才能以这种方式访问​​所有n索引？

标签： pythoncythonn-dimensional