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 术语，而不是以多进程方式并行化。

您的计算本质上是一个点(矩阵)积。对整个数组应用一次掩码，得到一个二维数组，NIJ。它的形状将是 (N,5)，其中 N 是 ~mask 中 True 值的数量。然后它只是一个 (5,N) 数组“点缀”了一个 (N,5) - 即。对 N 维度求和，得到一个 (5,5) 数组。

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

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

---

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

对于具有 400,000 个真值的 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() 只节省几毫秒。