加快python中的numpy循环？

Posted 2023-02-22

【中文标题】加快python中的numpy循环？

【英文标题】：Speeding up a numpy loop in python?

【发布时间】：2014-07-27 02:46:56

【问题描述】：

考虑以下使用非常慢的 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）

【问题讨论】：

您是否也考虑使用 OpenCL？

我不需要完整的性能，我只想要原始速度。

从点 tx, ty, tz 构建一个scipy.spatial.KDTree 然后对 ox, oy, oz 中的每个点使用无穷范数中的最近邻查找来查看是否有任何点足够接近。

你考虑过使用 Cython 吗？我读过它可以在没有太多痛苦的情况下提供很大的加速。 ***.com/questions/7799977/numpy-vs-cython-speed

【参考方案1】：

您正在分配 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  

【讨论】：

等一下，下一个来了:)

好的，矢量化形式在（可能需要检查条件），但是当你修改你的问题以回应我的原始答案时，部分答案变得无关紧要，即使它仍然是一个好的在非颠簸数据的情况下要做的初始步骤:)

【参考方案2】：

通过在 PyPy 下运行以下代码，您可能会获得良好的结果：http://pypy.org/（我们的 NumPy 集成说明在 https://bitbucket.org/pypy/numpy）

【参考方案3】：

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

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% 确定这是正确的。