Numpy：将掩码应用于值，然后取平均值，但并行

我有一个 1d numpy 值数组：

v = np.array([0, 1, 4, 0, 5])

此外，我有一个 2d numpy 布尔掩码数组（在生产中，有数百万个掩码）：

m = np.array([
    [True, True, False, False, False],
    [True, False, True, False, True],
    [True, True, True, True, True],
])

我想将掩码中的每一行应用到数组 v ，然后计算屏蔽值的平均值。

预期行为：

results = []
for mask in m:
    results.append(np.mean(v[mask]))

print(results) # [0.5, 3.0, 2.0]

很容易按顺序执行，但我确信有一个漂亮的并行版本？我发现的一个解决方案：

mask = np.ones(m.shape)
mask[~m] = np.nan
np.nanmean(v * mask, axis=1) # [0.5, 3.0, 2.0]

是否还有另一种解决方案，也许使用 np.ma 模块？我正在寻找比我当前的两个解决方案更快的解决方案。

可以使用布尔索引和矩阵乘法的组合来实现这一点，而无需任何显式循环或掩码创建。以下是方法：

v = np.array([0, 1, 4, 0, 5])
m = np.array([
    [True, True, False, False, False],
    [True, False, True, False, True],
    [True, True, True, True, True],
])

# 计算每个掩码的 True 值的数量
counts = m.sum(axis=1, keepdims=True)

# 使用布尔索引执行矩阵乘法以获取每个掩码的总和
sums = (m * v).sum(axis=1, keepdims=True)

# 计算平均值
results = sums / counts

print(results) # [[0.5], [3. ], [2. ]]

让我们分解一下：

1. counts = m.sum(axis=1, keepdims=True) :
2. 这会计算 m 中每行的 True 值的数量（即，每个掩码中的有效元素数量）。
3. axis=1 指定我们对行进行求和。

4. keepdims=True 保留了维度，从而使 counts 成为一个列向量，便于广播。

5. sums = (m * v).sum(axis=1, keepdims=True) :

6. m * v 使用广播将 v 乘以 m 中的每一行。由于 m 是布尔值，因此这实际上会选择 v 中对应于 True 值的元素。

7. (m * v).sum(axis=1, keepdims=True) 对每行求和，得到每个掩码的有效值之和，同时保持维度以实现广播。

8. results = sums / counts :

9. 最后，我们将每个掩码的总和除以每个掩码的计数，以获得平均值。广播用于按元素执行此操作。

这种方法利用了 NumPy 的矢量化操作和布尔索引，使其非常高效。它避免了显式循环和创建临时数组，例如的 nan 解决方案，使其在性能和内存使用方面可能更快。