如何加速Python中的numpy循环？

Question

考虑一下下面这段使用numpy数组的代码，它运行得非常慢：

# Intersection of an octree and a trajectory
def intersection(octree, trajectory):
    # Initialize numpy arrays
    ox = octree.get("x")
    oy = octree.get("y")
    oz = octree.get("z")
    oe = octree.get("extent")/2
    tx = trajectory.get("x")
    ty = trajectory.get("y")
    tz = trajectory.get("z")
    result = np.zeros(np.size(ox))
    # Loop over elements
    for i in range(0, np.size(tx)):
        for j in range(0, np.size(ox)):
            if (tx[i] > ox[j]-oe[j] and 
                tx[i] < ox[j]+oe[j] and 
                ty[i] > oy[j]-oe[j] and 
                ty[i] < oy[j]+oe[j] and 
                tz[i] > oz[j]-oe[j] and 
                tz[i] < oz[j]+oe[j]):
                result[j] += 1
    # Finalize
    return result

如何重写这个函数来加快计算速度呢？（np.size(tx) == 10000 和 np.size(ox) == 100000）

Answer 1

你可以尝试在PyPy这个环境下运行这段代码，通常会得到不错的效果。PyPy的官网是http://pypy.org/，关于如何让它和NumPy配合使用的说明可以在这里找到：https://bitbucket.org/pypy/numpy。

Answer 2

我觉得这样做可以让双重循环得到相同的结果，而且速度会更快：

for j in xrange(np.size(ox)):
    result[j] += sum( abs(tx-ox[j])<oe[j] & abs(ty-oy[j])<oe[j] & abs(tz-oz[j])<oe[j] )

要实现这个效果，首先：1) 重新排列循环的顺序（也就是交换它们的位置），这样做是可以的，因为在循环内部的内容没有变化；2) 把 result[j] 移到 i 循环外面；3) 将所有的 t>ox-oe 和 t<ox+oe 转换为 abs(t-ox)<oe（虽然这样可能不会大幅提升速度，但会更容易理解）。

由于你没有提供可以运行的代码，而我也不想为此构建一个测试，所以我不能百分之百确定这样做是正确的。

Answer 3

你正在分配10000个大小为100000的列表。首先要做的就是停止在内层的j循环中使用range，而是用生成器版本的xrange。这样可以节省分配这些列表时的时间和空间。

接下来，可以使用向量化操作：

for i in xrange(0, np.size(tx)):
    index = (ox-oe < tx[i]) & (ox+oe > tx[i]) & (oy-oe < ty[i]) & (oy+oe > ty[i]) & (oz-oe < tz[i]) & (oz+oe > tz[i])
    result[index] += 1