Numpy:改进对阵列的花哨索引

时间:2016-08-19 18:53:14

标签: python numpy optimization slice

为numpy寻找更快的花式索引,我正在运行的代码在np.take()处慢下来。我用order=F/C尝试np.reshape(),没有任何改进。 Python operator在没有双transpose的情况下效果很好,但是它们等于np.take().

p    = np.random.randn(3500, 51)
rows = np.asarray(range(p.shape[0]))
cols = np.asarray([1,2,3,4,5,6,7,8,9,10,15,20,25,30,40,50])

%timeit p[rows][:, cols]
%timeit p.take(cols, axis = 1 )
%timeit np.asarray(operator.itemgetter(*cols)(p.T)).T

1000 loops, best of 3: 301 µs per loop
10000 loops, best of 3: 132 µs per loop
10000 loops, best of 3: 135 µs per loop

1 个答案:

答案 0 :(得分:2)

对几个选项的测试:

In [3]: p[rows][:,cols].shape
Out[3]: (3500, 16)
In [4]: p[rows[:,None],cols].shape
Out[4]: (3500, 16)
In [5]: p[:,cols].shape
Out[5]: (3500, 16)
In [6]: p.take(cols,axis=1).shape
Out[6]: (3500, 16)

时间测试 - 普通p[:,cols]是最快的。尽可能使用切片。

In [7]: timeit p[rows][:,cols].shape
100 loops, best of 3: 2.78 ms per loop
In [8]: timeit p.take(cols,axis=1).shape
1000 loops, best of 3: 739 µs per loop
In [9]: timeit p[rows[:,None],cols].shape
1000 loops, best of 3: 1.43 ms per loop
In [10]: timeit p[:,cols].shape
1000 loops, best of 3: 649 µs per loop

我看到itemgetter用于列表,但不是数组。它是一个迭代一组索引的类。这两行正在做同样的事情:

In [23]: timeit np.asarray(operator.itemgetter(*cols)(p.T)).T.shape
1000 loops, best of 3: 738 µs per loop
In [24]: timeit np.array([p.T[c] for c in cols]).T.shape
1000 loops, best of 3: 748 µs per loop

请注意p.T[c]p.T[c,:]p[:,c].T。由于cols相对较少,并且忽略了使用rows的高级索引,因此它接近p[:,cols]