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

我想你想 'vectorize'，使用 numpy 术语，而不是以多进程方式并行化。

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

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

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

应用大布尔值 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() 只会节省几毫秒。

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 的数据类型）。

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

Python - parallelize a python loop for 2D masked array?

parallel-processing

loops

numpy

mask