如何将3D Numpy数组矢量化

Question

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

现在我正在使用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内置函数和矢量化函数处理相似数量元素的其他函数慢了约两个数量级）。

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

- ButterDog

---

你的代码中 d_n 是什么意思？ - Jaime

这只是一个与 b 形状相同的参数，只是修改了示例。 - ButterDog

Answer 1

在3D图像中通常的做法是将Z轴与第一个索引进行交换：

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

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

>>> for slice in a:
       do something

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

奖励：如果您连续存储数据，而不进行转置（或使用 a = np.ascontiguousarray(a.transpose((2,0,1))) 转换为连续数组），则访问 2D 切片将更快，因为它们在内存中连续映射。

- Imanol Luengo

---

不完全是我想要的答案，但我得到了更快的执行和更短的代码 :) - ButterDog

Answer 2

显然，你想要摆脱明确的 for 循环，但我认为这取决于你在 C 中使用的计算方式。以一个简单的例子为例，

a = np.zeros((100,100, 20))
a[:,:] = np.linspace(1,20,20)    # example data: 1,2,3,.., 20 as "z" for every "x","y"

b = np.sum(a[:,:]**2, axis=2)

将用a的平方和"z"值的总和填充100乘100数组b，即1+4+9+...+400 = 2870。

- xnx

---

这是一个相当复杂的函数，我不想包含它以避免污染示例，但基本上需要迭代C并找到一些最小值。整个过程可以使用numpy标准函数完成。 - ButterDog

你能否像我的例子一样用a[:, :]替换C？ - xnx

Answer 3

如果您的内部计算足够复杂，并且不适合矢量化，则您的迭代结构很好，并且不会对计算时间产生重大影响。

for (x,y), v in np.ndenumerate(b):
    C = a[x,y,:]
    ...
    for d in range(len(C)):
        ... # complex, not vectorizable calc
    ...
    b[x,y] = some_value

在第一和第二维中似乎没有特殊的结构，因此您可以将其视为 2D 映射到 1D，例如将 (N,20) 数组映射到一个 (N,) 数组。这并不会加速任何事情，但可以帮助突出问题的基本结构。

一步是专注于加快 C 到 some_value 的计算速度。有像 cumsum 和 cumprod 这样的函数，可以帮助您对向量进行顺序计算。 cython 也是一个很好的工具。

另一种方法是看看是否可以一次性处理那些内部计算的 N 个值。换句话说，如果必须迭代，最好是在最小的维度上进行。

从某种意义上说，这是一个非答案。但是，如果没有关于如何从 C 和 d_n 获取 some_value 的全部知识，我认为我们不能做更多。

---

看起来可以一次性计算所有点的e：

e = 2.5 * float(math.pow(d_n[x,y] - d, 2)) + C[d] * 0.05

E = 2.5 * (d_n[...,None] - np.arange(a.shape[-1]))**2 + a * 0.05  # (100,100,20)

E.min(axis=-1)  # smallest value along the last dimension
E.argmin(axis=-1)  # index of where that min occurs

乍一看，似乎这个E.argmin是您想要的b值（如果需要，可以针对某些边界条件进行调整）。

我没有真实的a和d_n数组，但是通过简单的测试，这个E.argmin(-1)与您的b匹配，速度提高了66倍。

- hpaulj

Answer 4

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

在NumPy中，许多函数是"缩减"函数，例如sum、any、std等。如果你给这些函数提供一个除None以外的axis参数，它将沿着该轴减少数组的维度。对于你的代码，你可以使用argmin函数，如果你首先以向量化的方式计算e。

d = np.arange(a.shape[2])
e = 2.5 * (d_n[...,None] - d)**2 + a*0.05
b = np.argmin(e, axis=2)

使用[...,None]进行索引是为了启用广播。 e中的值是浮点数，因此与sys.maxint进行比较有点奇怪，但是这就是它。

I, J = np.indices(b.shape)
b[e[I,J,b] >= sys.maxint] = a.shape[2] - 1

_{* 严格来说，一个约简函数的形式为reduce(operator, sequence)，因此严格地说不是std和argmin}

- user2379410