python - numpy 中的并行数组操作

Question

我有一段代码需要处理一些大的 numpy 数组。例如，我有一个 3D 数组 A，我需要使用 A 的元素构造另一个 3D 数组 B。然而，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,:])

如果每行并行运行，这也会加快速度。怎么办？

Answer 1

您应该查看Numpy's 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]