在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]