在numpy中更新整数索引切片的快速方法

时间:2014-06-21 09:30:33

标签: python optimization numpy

我有一个相当大的numpy数组,我正在进行清理数据(这是伪代码):

arr = np.ones(shape=(int(5E5), 1000), dtype=float)
to_zero = np.arange(500, 1000, 2) # Normally this is a function that finds 
                                  # columns be zeroed indexes are not continous
arr[:, to_zero] = 0

问题是arr[:, to_zero] = 0需要很长时间。在这个例子中,它需要4秒。 虽然arr[:, :500]需要500毫秒。

有没有办法让它更快?

numpy 1.8.1numpy 1.9 beta的numpy笔记本示例(请参阅时序结果)。

@Jaime指出,从长远来看,使用较新的numpy将是一个不错的选择。

2 个答案:

答案 0 :(得分:1)

从numpy的内部存储角度来看,重置大量连续内存的速度非常快。以你的例子为例,问题是你像疯子一样在记忆中跳跃,因为你正在对列进行归零。

我尝试反过来转动你的阵列,然后速度几乎提高了3倍。我的版本为4.08秒,转置版本为1.57秒。所以,如果你能做到这一点,那至少是一个很好的优化。

使用此索引可能会出现一些可疑的问题,因为实际上正在执行:

for c in to_zero:
    arr[:, c] = 0

比使用列表表示法更快。

所以,我跑了几个不同的选择:

to_zero = numpy.arange(0, 500, 2)

# 1. table in the original orientation
arr = numpy.ones((500000, 1000), dtype='float')

# 1.1. basic case: 4.08 s
arr[:, to_zero] = 0

# 1.2. a bit different notation: 234 ms 
arr[:, 0:500:2] = 0

# 1.3. for loop: 2.75 s
for c in to_zero:
    arr[:, c] = 0

# 2. transposed table
arr = numpy.ones((1000, 500000), dtype='float')

# 2.1. basic case: 1.47 s
arr[to_zero,:] = 0

# 2.2. a bit different notation: 105 ms 
arr[0:500:2,:] = 0

# 2.3. for loop: 112 ms
for r in to_zero:
    arr[r, :] = 0

这些已经使用IPython%timeit计时,因此结果可能会在一次运行到另一次运行时有所不同,但似乎有一种模式。转置你的桌子并使用循环。

答案 1 :(得分:0)

问题在于如何在内存中分配数据。

但是要解决您的问题,请将行读取为列,将列读取为行:

arr = np.ones(shape=(1000, int(5E5)), dtype=float)
to_zero = np.arange(500, 1000, 2) # Normally this is a function that finds 
                              # columns be zeroed indexes are not continous
arr[to_zero, :] = 0