向量化 for 循环 NumPy

我对 Python 比较陌生，并且有一个嵌套的 for 循环。由于 for 循环需要一段时间才能运行，因此我试图找到一种方法来向量化此代码，以便它可以运行得更快。

在本例中，coord 是一个 3 维数组，其中 coord[x, 0, 0] 和 coord[x, 0, 1] 是整数，coord[x, 0, 2] 是 0 或 1。 H 是 SciPy稀疏矩阵，x_dist、y_dist、z_dist 和 a 都是浮点数。

# x_dist, y_dist, and z_dist are floats
# coord is a num x 1 x 3 numpy array where num can go into the hundreds of thousands
num = coord.shape[0]    
H = sparse.lil_matrix((num, num))
for i in xrange(num):
    for j in xrange(num):
        if (np.absolute(coord[i, 0, 0] - coord[j, 0, 0]) <= 2 and
                (np.absolute(coord[i, 0, 1] - coord[j, 0, 1]) <= 1)):

            x = ((coord[i, 0, 0] * x_dist + coord[i, 0, 2] * z_dist) -
                 (coord[j, 0, 0] * x_dist + coord[j, 0, 2] * z_dist))

            y = (coord[i, 0, 1] * y_dist) - (coord[j, 0, 1] * y_dist)

            if a - 0.5 <= np.sqrt(x ** 2 + y ** 2) <= a + 0.5:
                H[i, j] = -2.7

我还读到，使用 NumPy 进行广播虽然速度更快，但可能会导致临时数组使用大量内存。走矢量化路线或尝试使用 Cython 之类的东西会更好吗？

这就是我对代码进行矢量化的方式，稍后对注意事项进行一些讨论：

import numpy as np
import scipy.sparse as sps

idx = ((np.abs(coord[:, 0, 0] - coord[:, 0, 0, None]) <= 2) &
       (np.abs(coord[:, 0, 1] - coord[:, 0, 1, None]) <= 1))

rows, cols = np.nonzero(idx)
x = ((coord[rows, 0, 0]-coord[cols, 0, 0]) * x_dist +
     (coord[rows, 0, 2]-coord[cols, 0, 2]) * z_dist)
y = (coord[rows, 0, 1]-coord[cols, 0, 1]) * y_dist
r2 = x*x + y*y

idx = ((a - 0.5)**2 <= r2) & (r2 <= (a + 0.5)**2)

rows, cols = rows[idx], cols[idx]
data = np.repeat(2.7, len(rows))

H = sps.coo_matrix((data, (rows, cols)), shape=(num, num)).tolil()

正如您所指出的，问题将出现在第一个idx数组，因为它的形状(num, num)，所以如果num就是“入数十万”。

一种可能的解决方案是将问题分解为可管理的部分。如果您有一个包含 100,000 个元素的数组，则可以将其拆分为 100 个块，每块包含 1,000 个元素，并对 10,000 个块组合中的每一个运行上述代码的修改版本。您只需要 1,000,000 个元素idx数组（您可以预先分配和重用以获得更好的性能），并且您将拥有一个仅 10,000 次迭代的循环，而不是当前实现的 10,000,000,000 次迭代。这是一种穷人的并行化方案，如果您有一台多核机器，您实际上可以通过并行处理其中几个块来改进该方案。