使用掩码索引对NumPy数组进行向量化求和操作

Question

我将尝试使用numpy数组进行带有掩码索引的向量化求和操作。

例如，不使用掩码：

import numpy as np

# data to be used in a vectorized sum operation
data = np.array([[1,0,0,0,0,0],
                 [0,1,0,0,0,0],
                 [0,0,1,0,0,0]])

# data indices i wish to sum together
idx = np.array([[0,1,2],   # sum data rows 0,1 and 2
                [2,1,1]])  # sum data rows 2,1 and 1

# without a mask this is straighforward
print np.sum(data[idx],axis=1)
#[[1 1 1 0 0 0]
# [0 2 1 0 0 0]]

现在有了掩码，我无法找出如何在不循环遍历掩码索引数组的情况下完成它：

# introduce a mask
mask = np.array([[True,  True, True],  # sum data rows 0,1 and 2
                 [False, True, True]]) # sum data rows 1 and 1 (masking out idx[1,0])

summed = np.zeros((idx.shape[0],data.shape[1]),dtype='int')
for i in xrange(idx.shape[0]):
    summed[i] =  np.sum(data[idx[i][mask[i]]],axis=0)
print summed
#[[1 1 1 0 0 0]
 #[0 2 0 0 0 0]]

问题

在没有循环的情况下，有没有适当的方法执行这种类型的操作？

- Fnord

---

你的问题不再有意义了。掩码的第二行对数据的第1行和第2行求和，而不是对第1行和第1行求和。你无法从掩码中获得重复的元素... - Mad Physicist

Answer 1

您可以使用 np.einsum 来解决这个问题。

v = data[idx]
summed = np.einsum('ijk,ij->ik', v, mask)

运行给定的样例 -

In [43]: v = data[idx]

In [44]: np.einsum('ijk,ij->ik', v, mask)
Out[44]: 
array([[1, 1, 1, 0, 0, 0],
       [0, 2, 0, 0, 0, 0]])

或者，使用np.matmul函数 -

In [67]: np.matmul(v.swapaxes(1,2), mask[...,None])[...,0]
Out[67]: 
array([[1, 1, 1, 0, 0, 0],
       [0, 2, 0, 0, 0, 0]])

# Put another way
In [80]: np.matmul(mask[:,None,:], v)[:,0]
Out[80]: 
array([[1, 1, 1, 0, 0, 0],
       [0, 2, 0, 0, 0, 0]])

---

保持循环并提高性能

如果你没有足够的循环次数，并且每次迭代都存在足够的求和操作，那么可以将迭代操作替换为矩阵乘法操作。因此 -

for i in xrange(idx.shape[0]):
    summed[i] = mask[i].dot(data[idx[i]])

- Divakar