快速的numpy花式索引

Question

我的切片numpy数组的代码(通过花哨的索引)非常慢.它目前是计划的瓶颈.

a.shape
(3218, 6)

ts = time.time(); a[rows][:, cols]; te = time.time(); print('%.8f' % (te-ts));
0.00200009

获取由矩阵a的行'rows'和列'col'的子集组成的数组的正确numpy调用是什么？(事实上​​,我需要这个结果的转置)

Answer 1

让我尝试总结Jaime和TheodrosZelleke的优秀答案,并在一些评论中加入.

1. 高级(花哨)索引始终返回副本,而不是视图.
2. a[rows][:,cols]意味着两个花哨的索引操作,因此a[rows]创建并丢弃中间副本.方便可读,但效率不高.此外要注意,[:,cols]通常会生成一个C-cont的Fortran连续副本.资源.
3. a[rows.reshape(-1,1),cols]是一个基于事实的单一高级索引表达式,rows.reshape(-1,1)并且cols被广播到预期结果的形状.

4. 一个常见的经验是,扁平数组中的索引可能比花式索引更有效,因此另一种方法是

   indx = rows.reshape(-1,1)*a.shape[1] + cols
   a.take(indx)
   

   要么

   a.take(indx.flat).reshape(rows.size,cols.size)
   

5. 效率将取决于内存访问模式以及起始数组是C-countinous还是Fortran连续,因此需要进行实验.

6. 仅在真正需要时使用花式索引:基本切片 a[rstart:rstop:rstep, cstart:cstop:cstep]返回视图(虽然不是连续的)并且应该更快!

Answer 2

令我惊讶的是,计算第一个线性1D指数的长度表达式比问题中提出的连续数组索引快50%以上:

(a.ravel()[(
   cols + (rows * a.shape[1]).reshape((-1,1))
   ).ravel()]).reshape(rows.size, cols.size)

更新: OP更新了初始数组形状的描述.随着更新的大小,加速现在超过99%:

In [93]: a = np.random.randn(3218, 1415)

In [94]: rows = np.random.randint(a.shape[0], size=2000)

In [95]: cols = np.random.randint(a.shape[1], size=6)

In [96]: timeit a[rows][:, cols]
10 loops, best of 3: 186 ms per loop

In [97]: timeit (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size)
1000 loops, best of 3: 1.56 ms per loop

INITAL ANSWER: 以下是成绩单:

In [79]: a = np.random.randn(3218, 6)
In [80]: a.shape
Out[80]: (3218, 6)

In [81]: rows = np.random.randint(a.shape[0], size=2000)
In [82]: cols = np.array([1,3,4,5])

时间方法1:

In [83]: timeit a[rows][:, cols]
1000 loops, best of 3: 1.26 ms per loop

时间方法2:

In [84]: timeit (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size)
1000 loops, best of 3: 568 us per loop

检查结果是否实际相同:

In [85]: result1 = a[rows][:, cols]
In [86]: result2 = (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size)

In [87]: np.sum(result1 - result2)
Out[87]: 0.0

Answer 3

如果使用花式索引和广播切片,您可以加快速度:

from __future__ import division
import numpy as np

def slice_1(a, rs, cs) :
    return a[rs][:, cs]

def slice_2(a, rs, cs) :
    return a[rs[:, None], cs]

>>> rows, cols = 3218, 6
>>> rs = np.unique(np.random.randint(0, rows, size=(rows//2,)))
>>> cs = np.unique(np.random.randint(0, cols, size=(cols//2,)))
>>> a = np.random.rand(rows, cols)
>>> import timeit
>>> print timeit.timeit('slice_1(a, rs, cs)',
                        'from __main__ import slice_1, a, rs, cs',
                        number=1000)
0.24083110865
>>> print timeit.timeit('slice_2(a, rs, cs)',
                        'from __main__ import slice_2, a, rs, cs',
                        number=1000)
0.206566124519

如果按照百分比来考虑,做一些快15%的事情总是好的,但是在我的系统中,对于阵列的大小,这需要40美元来做切片,并且很难相信操作采取240我们将成为你的瓶颈.