在python中加速numpy循环?

时间:2014-06-05 22:54:45

标签: python loops optimization numpy iteration

使用非常慢的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) == 10000np.size(ox) == 100000

3 个答案:

答案 0 :(得分:6)

您正在分配10000个大小为100000的列表。首先要做的是停止对嵌套range循环使用j并使用生成器版本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  

答案 1 :(得分:0)

您可以通过在PyPy下运行此代码来获得良好的结果:http://pypy.org/(我们在https://bitbucket.org/pypy/numpy进行NumPy集成的说明)

答案 2 :(得分:0)

我认为这将为双循环提供相同的结果并且更快:

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 and t<ox+oe转换为abs(t-ox)<oe(虽然这可能不是一个巨大的加速,但它更容易阅读)。

由于您没有可运行的代码,并且我不想为此构建测试,因此我不能100%确定这是正确的。