Быстрое (er) причудливое индексирование и сокращение?

Я пытаюсь использовать и ускорить причудливую индексацию, чтобы «объединить» два массива и суммировать по одной из осей результатов.

Что-то вроде этого:

$ ipython
In [1]: import numpy as np
In [2]: ne, ds = 12, 6
In [3]: i = np.random.randn(ne, ds).astype('float32')
In [4]: t = np.random.randint(0, ds, size=(1e5, ne)).astype('uint8')

In [5]: %timeit i[np.arange(ne), t].sum(-1)
10 loops, best of 3: 44 ms per loop

Есть ли простой способ ускорить оператор в In [5]? Должен ли я использовать OpenMP и что-то вроде scipy.weave или prange Cython?


python numpy scipy optimization cython
person npinto    schedule 03.08.2012    source источник
comment
Другой связанный с этим вопрос: как бы я использовал pandas, чтобы сделать то же самое?   -  person npinto    schedule 03.08.2012
comment
Numpy делает это на скорости C, поэтому вы, вероятно, не сможете сильно ускорить его с помощью weave.   -  person reptilicus    schedule 03.08.2012

Ответы (1)


arrow_upward
8
arrow_downward

numpy.take по какой-то причине намного быстрее, чем причудливое индексирование. Единственный трюк заключается в том, что он рассматривает массив как плоский.

In [1]: a = np.random.randn(12,6).astype(np.float32)

In [2]: c = np.random.randint(0,6,size=(1e5,12)).astype(np.uint8)

In [3]: r = np.arange(12)

In [4]: %timeit a[r,c].sum(-1)
10 loops, best of 3: 46.7 ms per loop

In [5]: rr, cc = np.broadcast_arrays(r,c)

In [6]: flat_index = rr*a.shape[1] + cc

In [7]: %timeit a.take(flat_index).sum(-1)
100 loops, best of 3: 5.5 ms per loop

In [8]: (a.take(flat_index).sum(-1) == a[r,c].sum(-1)).all()
Out[8]: True

Я думаю, что единственный другой способ, которым вы увидите значительное улучшение скорости, — это написать собственное ядро ​​для графического процессора, используя что-то вроде PyCUDA.

person user545424    schedule 05.08.2012
comment
По умолчанию массив обрабатывается только как плоский, вы все равно можете использовать ключевое слово axis. То есть np.take(np.arange(10).reshape((-1,2)), [0], axis=0) выберет первую строку. - person jorgeca; 07.08.2012
comment
@jorgeca: правильно, но я не думаю, что вы можете вытащить отдельные элементы, указав и строку, и столбец, как вы можете с причудливой индексацией, если вы не индексируете плоский массив. - person user545424; 07.08.2012