Python - 为 2D 掩码数组并行化 Python 循环？

Question

可能是一个常见的问题，但如何在 Python 中并行化这个循环？

for i in range(0,Nx.shape[2]):
  for j in range(0,Nx.shape[2]):
    NI=Nx[:,:,i]; NJ=Nx[:,:,j]
    Ku[i,j] = (NI[mask!=True]*NJ[mask!=True]).sum()

所以我的问题是：并行化这段代码最简单的方法是什么？

         ---------- EDIT LATER------------------

数据示例

import random
import numpy as np
import numpy.ma as ma
from numpy import unravel_index    

#my input
Nx = np.random.rand(5,5,5)  

#mask creation
mask_positions = zip(*np.where((Nx[:,:,0] < 0.4)))
mask_array_positions = np.asarray(mask_positions)
i, j = mask_array_positions.T
mask = np.zeros(Nx[:,:,0].shape, bool)
mask[i,j] = True

我想通过并行计算 Ku。我的目标是使用 Ku 数组来解决线性问题，所以我必须将掩码值分开（代表我数组的一半）

Answer 1

我认为您想“向量化”，使用numpy 术语，而不是以多进程方式并行化。

您的计算本质上是一个点（矩阵）积。将mask 一次应用于整个数组以获得二维数组NIJ。它的形状将是(N,5)，其中N 是~mask 中True 值的数量。然后它只是一个带有(N,5) 的(5,N) 数组'点缀' - 即。对N 维度求和，留下一个(5,5) 数组。

NIJ = Nx[~mask,:]
Ku = np.dot(NIJ.T,NIJ)

在快速测试中，它与双循环产生的Ku 匹配。根据np.dot 使用的底层库，可能会有一些多核计算，但这通常不是numpy 用户的优先问题。

---

应用大布尔值 mask 是这些计算中最耗时的部分 - 无论是矢量化版本还是迭代版本。

对于具有 400,000 个 True 值的 mask，比较这两个索引时间：

In [195]: timeit (NI[:400,:1000],NJ[:400,:1000])
100000 loops, best of 3: 4.87 us per loop
In [196]: timeit (NI[mask],NJ[mask])
10 loops, best of 3: 98.8 ms per loop

使用基本（切片）索引选择相同数量的项目比使用 mask 的高级索引快几个数量级。

用np.dot(NI[mask],NJ[mask]) 替换(NI[mask]*NJ[mask]).sum() 只会节省几毫秒。

Answer 2

我想将 @hpaulj 的出色答案（顺便说一下，对问题的出色分析）扩展到大型矩阵。

操作

Ku = np.dot(NIJ.T,NIJ)

可以替换为

Ku = np.einsum('ij,ik->jk', NIJ, NIJ)

还应注意np.dot could fall back to slower routines 如果 numpy 未编译为使用 BLAS。

对于形状为(1250711, 50) 的测试矩阵NIJ，我使用dot 方法得到54.9 s，而einsum 在1.67 s 中得到它。在我的系统上，numpy 编译时支持 BLAS。

备注：np.einsum 并不总是优于 np.dot，当您比较以下任何一项时，这种情况在我的系统上变得很明显

Nx = np.random.rand(1400,1528,20).astype(np.float16)
Nx = np.random.rand(1400,1528,20).astype(np.float32)

（甚至是np.float64的dtype）。