Итерация по произвольной размерности numpy.array

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

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

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(axis = 0).

Ответ 1

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

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

Ответ 2

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

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

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

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

Ответ 3

Таким образом, можно легко перебрать первое измерение, как вы показали. Другой способ сделать это для произвольного измерения - это использовать 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)

ОБНОВЛЕНИЕ: Я прокомментирую, что я отправил PR для numpy, чтобы решить эту проблему здесь: https://github.com/numpy/numpy/pull/3262. Концепция заключалась в том, что этого было недостаточно, чтобы добавить к единой кодовой базе. Я думаю, что использование np.rollaxis - лучший способ сделать это, и, если вам нужен интегратор, поместите его в iter().

Ответ 4

Я думаю, что нет никакой функции. Когда я написал свою функцию, я закончил тем, что предложил итерацию 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,)]

Ответ 5

Вы можете использовать numpy.shape, чтобы получить размеры, а затем выполнить диапазон для их итерации.

n0, n1, n2 = numpy.shape(c)

for r in range(n0):
    print(c[r,:,:])