Итерация по произвольному измерению в NumPy

Question

У меня есть многомерный массив NumPy, и мне нужно выполнить итерацию по заданному измерению.Проблема в том, что я не буду знать , какое измерение до времени выполнения.Другими словами, учитывая массив m, я мог бы хотеть

m[:,:,:,i] for i in xrange(n)

или я мог бы хотеть

m[:,:,i,:] for i in xrange(n)

и т. Д.

Я предполагаю, что должно быть простовозможность написать это, но я не могу понять, что это такое / как это можно назвать.Есть мысли?

Answer 1

Есть много способов сделать это. Вы можете создать правильный индекс со списком срезов или, возможно, изменить шаги m. Однако самым простым способом может быть использование np.swapaxes:

import numpy as np
m=np.arange(24).reshape(2,3,4)
print(m.shape)
# (2, 3, 4)

Пусть axis будет осью, которую вы хотите зациклить. m_swapped совпадает с m, за исключением того, что ось axis=1 поменялась местами с последней (axis=-1) осью.

axis=1
m_swapped=m.swapaxes(axis,-1)
print(m_swapped.shape)
# (2, 4, 3)

Теперь вы можете просто пройти по последней оси:

for i in xrange(m_swapped.shape[-1]):
    assert np.all(m[:,i,:] == m_swapped[...,i])

Обратите внимание, что m_swapped является представлением, а не копией m. Изменение m_swapped изменит m.

m_swapped[1,2,0]=100
print(m)
assert(m[1,0,2]==100)

Answer 2

Вы можете использовать slice(None) вместо :.Например,

from numpy import *

d = 2  # the dimension to iterate

x = arange(5*5*5).reshape((5,5,5))
s = slice(None)  # :

for i in range(5):
    slicer = [s]*3  # [:, :, :]
    slicer[d] = i   # [:, :, i]
    print x[slicer] # x[:, :, i]