numpy中的并行数组操作

Question

我有一段代码，需要处理一些很大的numpy数组。比如说，我有一个三维数组A，我想用A里的元素来构建另一个三维数组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 中进行并行处理，可以看看 mpi4py。这是一个为Python提供的MPI接口，它可以让你进行分布式处理。

Answer 2

如果你已经好好处理了边缘部分，那么进行第一次向量化的标准方法大概是这样的：

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] 填充合适的值。

把这个方法扩展到其他索引应该也很简单。

Answer 3

你可以看看 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]