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是True中~mask值的数量。然后它只是一个(5,N)数组'点''与(N,5) - 即。总结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。

对于形状NIJ的测试矩阵(1250711, 50)，我使用54.9 s方法得到dot，而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。

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

2 个答案: