循环NumPy的矢量化

Question

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

在这种情况下，coord是一个三维数组，其中coord[x，0，0]和coord[x，0，1]是整数，coord[y，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的东西更好？

Answer 1

这就是我将如何矢量化您的代码，稍后将讨论一些注意事项：

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"达到数十万"，可能会把您的记忆炸成碎片

一个潜在的解决方案是将你的问题分解成可管理的部分。如果您有一个100000元素的数组，您可以将其拆分为100个1000元素的块，并为10000个块的组合中的每一个运行上面代码的修改版本。您只需要一个1000000元素的idx数组（您可以预先分配并重用它以获得更好的性能），并且您将拥有一个只有10000次迭代的循环，而不是当前实现的10000000000次。这是一种穷人的并行化方案，如果你有一台多核机器，你可以通过并行处理其中的几个块来改进它。

Answer 2

计算的性质使我很难用我熟悉的numpy方法进行矢量化。我认为就速度和内存使用而言，最好的解决方案是cython。然而，使用numba可以获得一些加速。这里有一个例子（注意，通常使用autojit作为装饰器）：

import numpy as np
from scipy import sparse
import time
from numba.decorators import autojit
x_dist=.5
y_dist = .5
z_dist = .4
a = .6
coord = np.random.normal(size=(1000,1000,1000))
def run(coord, x_dist,y_dist, z_dist, a):
    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
    return H
runaj = autojit(run)
t0 = time.time()
run(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'First Original Runtime:', t1 - t0
t0 = time.time()
run(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Second Original Runtime:', t1 - t0
t0 = time.time()
run(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Third Original Runtime:', t1 - t0
t0 = time.time()
runaj(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'First Numba Runtime:', t1 - t0
t0 = time.time()
runaj(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Second Numba Runtime:', t1 - t0
t0 = time.time()
runaj(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Third Numba Runtime:', t1 - t0

我得到这个输出：

First Original Runtime: 21.3574919701
Second Original Runtime: 15.7615520954
Third Original Runtime: 15.3634860516
First Numba Runtime: 9.87108802795
Second Numba Runtime: 9.32944011688
Third Numba Runtime: 9.32300305367