Итерация по произвольному измерению в numpy
У меня есть многомерный массив Numpy, и мне нужно выполнить итерацию по заданному измерению. Проблема в том, что я не буду знать , какое измерение до времени выполнения. Другими словами, учитывая массив m, я мог бы хотеть
m[:,:,:,i] for i in xrange(n)
или я мог хотеть
m[:,:,i,:] for i in xrange(n)
и т. Д.
Я представляю, что в numpy должна быть простая функция, чтобы написать это, но я могу ' не понять, что это такое / как это можно назвать. Есть мысли?
2 ответа
Есть много способов сделать это. Вы можете построить правильный индекс со списком срезов или, возможно, изменить шаги m . Однако самым простым способом может быть использование np.swapaxes :
import numpy as np
m=np.arange(24).reshape(2,3,4)
print(m.shape)
# (2, 3, 4)
Пусть ось будет осью, которую вы хотите перебрать. 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)
Вы можете использовать срез (Нет) вместо : . Например,
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]