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是數量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)

還應注意，如果未將numpy編譯為使用BLAS，則np.dot 可能會退回到較慢的例程 。

對於形狀為(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 ）。