在numpy中遍历任意维数

Question

我有一个多维的numpy数组，我需要遍历给定的维度。问题是，我不知道哪个维直到运行。换句话说，给定一个数组男，我也想

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

或者我可以想

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

等

我想，必须有numpy的写这一个简单的功能，但我无法弄清楚它是什么/它可能被称为什么。有什么想法吗？

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]