2009-10-19 2 views
34

Есть ли функция для получения итератора по произвольной размерности массива numpy?Итерация по произвольной размерности numpy.array

Итерация по первому измерению легко ...

In [63]: c = numpy.arange(24).reshape(2,3,4) 

In [64]: for r in c : 
    ....:  print r 
    ....: 
[[ 0 1 2 3] 
[ 4 5 6 7] 
[ 8 9 10 11]] 
[[12 13 14 15] 
[16 17 18 19] 
[20 21 22 23]] 

Но итерацию по сравнению с другими размерами сложнее. Например, последнее измерение:

In [73]: for r in c.swapaxes(2,0).swapaxes(1,2) : 
    ....:  print r 
    ....: 
[[ 0 4 8] 
[12 16 20]] 
[[ 1 5 9] 
[13 17 21]] 
[[ 2 6 10] 
[14 18 22]] 
[[ 3 7 11] 
[15 19 23]] 

Я делаю генератор, чтобы сделать это сам, но я удивлен, что это не функция с именем что-то вроде numpy.ndarray.iterdim (ось = 0) в сделайте это автоматически.

ответ

36

Что вы предлагаете довольно быстро, но разборчивость может быть улучшена с более четкими формами:

for i in range(c.shape[-1]): 
    print c[:,:,i] 

или, лучше (быстрее, более общее и более четко):

for i in range(c.shape[-1]): 
    print c[...,i] 

Однако , первый подход выше, по-видимому, примерно в два раза медленнее, чем подход swapaxes():

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for r in c.swapaxes(2,0).swapaxes(1,2): u = r' 
100000 loops, best of 3: 3.69 usec per loop 

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for i in range(c.shape[-1]): u = c[:,:,i]' 
100000 loops, best of 3: 6.08 usec per loop 

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for r in numpy.rollaxis(c, 2): u = r' 
100000 loops, best of 3: 6.46 usec per loop 

Я бы предположил, что это потому, что swapaxes() не копирует никаких данных, и потому что обработка c[:,:,i] может быть выполнена с помощью общего кода (который обрабатывает случай, когда : заменен более сложным фрагментом).

Обратите внимание, однако, что более явное второе решение c[...,i] является одновременно вполне разборчиво и довольно быстро:

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for i in range(c.shape[-1]): u = c[...,i]' 
100000 loops, best of 3: 4.74 usec per loop 
3

Я предполагаю, что нет никакой функции. Когда я написал свою функцию, я закончил тем, что предложил итерацию EOL. Для будущих читателей, здесь:

def iterdim(a, axis=0) : 
    a = numpy.asarray(a); 
    leading_indices = (slice(None),)*axis 
    for i in xrange(a.shape[axis]) : 
    yield a[leading_indices+(i,)] 
+0

Стандартный синтаксис NumPy 'a [..., i]' будет легче и устранит необходимость в 'leading_indices'. – EOL

+1

@EOL, но это будет работать только для последней оси, а ведущие - более общие ... – lukas

+0

Хорошая точка @lukas: исходный вопрос действительно упоминает итерацию «над произвольным измерением», а я имел в виду интеграцию по последнему измерению. – EOL

14

Я хотел бы использовать следующее:

c = numpy.arange(2 * 3 * 4) 
c.shape = (2, 3, 4) 

for r in numpy.rollaxis(c, 2): 
    print(r) 

Функции создают поперечные новый вид на массиве. В этом случае он перемещает ось 2 вперед, что эквивалентно операции c.transpose(2, 0, 1).

+1

+1: Совсем прямо, но, к сожалению, чуть медленнее, чем простой подход 'c [:,:, i]' (не знаете почему). – EOL

4

Итак, как вы показали, можно легко итотировать в первом измерении. Другой способ сделать это для произвольной размерности - использовать numpy.rollaxis(), чтобы перенести заданное измерение в первое (поведение по умолчанию), а затем использовать возвращенный массив (который является представлением, так что это быстро) в качестве итератора ,

In [1]: array = numpy.arange(24).reshape(2,3,4) 

In [2]: for array_slice in np.rollaxis(array, 1): 
    ....:  print array_slice.shape 
    ....: 
(2, 4) 
(2, 4) 
(2, 4) 

EDIT: Я буду комментировать, что я представил PR в Numpy для решения этой проблемы здесь: https://github.com/numpy/numpy/pull/3262. Консенсус заключался в том, что этого недостаточно, чтобы добавить к базе данных numpy. Я думаю, что использование np.rollaxis - лучший способ сделать это, и если вы хотите сделать interator, оберните его в iter().