Мой код для нарезки массива numpy (через причудливую индексацию) очень медленный. В настоящее время это узкое место в программе.
a.shape
(3218, 6)
ts = time.time(); a[rows][:, cols]; te = time.time(); print('%.8f' % (te-ts));
0.00200009
Каков правильный вызов numpy для получения массива, состоящего из подмножества строк "rows" и столбцов "col" матрицы a? (на самом деле мне нужно транспонировать этот результат)
Вы можете получить некоторую скорость, если вы используете фрагментацию с помощью удобной индексации и трансляции:
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 us, станет вашим узким местом.
Позвольте мне попытаться обобщить отличные ответы Хайме и Теодроселеке и смешать некоторые комментарии.
a[rows][:,cols] подразумевает две причудливые операции индексирования, поэтому создается промежуточная копия a[rows] и отбрасывается. Удобный и читаемый, но не очень эффективный. Кроме того, остерегайтесь того, что [:,cols] обычно генерирует непрерывную копию Фортрана C-cont. источник.a[rows.reshape(-1,1),cols] - это одно расширенное выражение индексации, основанное на том, что rows.reshape(-1,1) и cols broadcast в форму от предполагаемого результата.Общим опытом является то, что индексирование в сплющенном массиве может быть более эффективным, чем причудливое индексирование, поэтому другой подход
indx = rows.reshape(-1,1)*a.shape[1] + cols
a.take(indx)
или
a.take(indx.flat).reshape(rows.size,cols.size)
Эффективность будет зависеть от шаблонов доступа к памяти и будет ли исходный массив C-countinous или Fortran непрерывным, поэтому необходимо провести эксперименты.
Использовать причудливую индексацию только если это действительно необходимо: базовая нарезка a[rstart:rstop:rstep, cstart:cstop:cstep] возвращает представление (хотя и не непрерывное) и должно быть быстрее
К моему удивлению, такое выражение длины, которое вычисляет первые линейные 1D-индексы, более чем 50% быстрее, чем последовательная индексация массива, представленная в вопросе:
(a.ravel()[(
cols + (rows * a.shape[1]).reshape((-1,1))
).ravel()]).reshape(rows.size, cols.size)
UPDATE: 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