如何矢量化 3D Numpy 数组

Question

我有一个 3D numpy 数组，如a = np.zeros((100,100, 20)). 我想对x,y涉及z轴上所有元素的每个位置执行一个操作，并将结果存储在一个数组中，就像b = np.zeros((100,100))在相同的相应x,y位置上一样。

现在我正在使用 for 循环：

d_n = np.array([...]) # a parameter with the same shape as b
for (x,y), v in np.ndenumerate(b):
    C = a[x,y,:]

    ### calculate some_value using C
    minv = sys.maxint
    depth = -1
    C = a[x,y,:]
    for d in range(len(C)):
        e = 2.5 * float(math.pow(d_n[x,y] - d, 2)) + C[d] * 0.05
        if e < minv:
            minv = e
            depth = d

    some_value = depth
    if depth == -1:
        some_value = len(C) - 1
    ###

    b[x,y] = some_value

现在的问题是这个操作比其他人用pythonic方式做的要慢得多，例如c = b * b（我实际上分析了这个函数，它比使用numpy内置函数和向量化函数的其他函数慢大约2个数量级，元素数量相似）

如何提高将 3D 数组映射到 2D 数组的此类函数的性能？

Answer 1

在 3D 图像中通常所做的是将Z轴交换到第一个索引：

>>> a = a.transpose((2,0,1))
>>> a.shape
(20, 100, 100)

现在您可以轻松地迭代 Z 轴：

>>> for slice in a:
       do something

这里将是3D 矩阵的slice每个部分。100x100此外，通过转置，您可以通过索引第一个轴来直接访问每个 2D 切片。例如，a[10]将为您提供第 11 个 2D100x100切片。

好处：如果您连续存储数据，而无需转置（或使用a = np.ascontiguousarray(a.transpose((2,0,1)))对 2D 切片的访问转换为连续数组），那么速度会更快，因为它们在内存中连续映射。