python - numpy 中的并行数组操作

Question

我有一个代码，我需要在其中处理一些大的 numpy 数组。例如，我有一个 3D 数组A，我需要B使用A. 然而，所有的元素B都是相互独立的。例子：

for i in np.arange(Nx):
  for j in np.arange(Ny):
   for k in np.arange(Nz):
       B[i][j][k] = A[i+1][j][k]*np.sqrt(A[i][j-1][k-1])

B因此，如果我可以并行构建数组，它将大大加快速度。在 python 中执行此操作的最简单方法是什么？

我也有类似的矩阵运算，比如标准化二维数组的每一行。例子

for i in np.arange(Nx):
   f[i,:] = f[i,:]/np.linalg.norm(f[i,:])

如果它为每一行并行运行，这也会加快速度。如何做呢？

score 2 · Accepted Answer

你应该看看Numpy 的roll功能。我认为这相当于您的第一个代码块（尽管您需要决定在边缘发生什么 - roll“环绕”）：

B = np.roll(A,1,axis=0) * np.sqrt(np.roll(np.roll(A,-1,axis=1),-1,axis=2))

对于您的第二种情况，另一个相当可怕的单线是：

f /= np.sqrt(np.sum(f**2, axis=1))[...,np.newaxis]

这条线的解释：

我们首先要计算每一行的范数。让我们

f = np.random.rand(5,6)

平方的每个元素f

f**2

将沿轴 1 的正方形相加，使该轴“变平”。

np.sum(f**2, axis=1)

取平方和的平方根。

np.sqrt(np.sum(f**2, axis=1))

我们现在有了每一行的范数。

为了正确地除以原始的每一行，f我们需要利用 Numpy 广播规则来有效地添加一个维度：

np.sqrt(np.sum(f**2, axis=1))[...,np.newaxis]

最后我们计算我们的结果

f /= np.sqrt(np.sum(f**2, axis=1))[...,np.newaxis]

score 0 · Accepted Answer

如果您对边缘进行了很好的处理，那么进行第一次矢量化的标准方法将是这样的：

B = np.zeros(A.shape)
B[:-1, 1:, 1:] = A[1:, 1:, 1:] * np.sqrt(A[:-1, :-1, :-1])

然后，您需要填写B[-1, :, :],B[:, 0, :]和B[:, :, 0]适当的值。

将此扩展到其他索引应该非常简单。

score 0 · Accepted Answer

要在中执行并行处理numpy，您应该查看mpi4py。它是 Python 的 MPI 绑定。它允许分布式处理。

3 回答 3