python - 如何用掩码和索引替换双 for 循环？

Question

我有几个嵌套的 for 循环可以做正确的事情（数组的屏蔽副本）。但是性能太慢了，我觉得必须有更好的 Pythonic 方式来做到这一点。目标是使用掩码确定何时从源复制数据，使用 coord 作为索引到源。有效的循环代码如下：

import numpy as np
dest = np.zeros((4,4,2))
source = range(32)
source = np.reshape(source,(4,4,2))
mask = np.ones((4,4),bool)
mask[1,0] = 0
coord = np.ones((4,4,2),int)

for y in range (0,dest.shape[0]):
    for x in range (0, dest.shape[1]):
        if np.all(mask[y,x]):
            dest[y,x] = source[coord[y,x,0], coord[y,x,1]]

print dest

运行后dest是这样的：

[[[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[  0.   0.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]]

source[1,1]被复制到所有的dest，除了dest[1,0]因为mask[1,0]设置为False。其余的mask是True. 谁能告诉我如何用更有效的东西替换循环？

Answer 1

使用numpy.where。您必须添加一个额外的维度，mask以便它会广播。

dest = np.where(mask[:,:,None], source[coord[:,:,0], coord[:,:,1]], dest)

或者如果合适的话：

dest = np.where(mask[:,:,None], source[coord[:,:,0], coord[:,:,1]], np.zeros((4,4,2)))