從多維數組中選擇一些維數

Question

我有一個4維數組，我需要從給定維數的這個數組中處理所有1維向量。這種運作良好：

 
    def myfun(arr4d,selected_dim): # selected_dim can only be 2 or 3 
     print arr4d.shape # (2, 5, 10, 10) 
     for i in xrange(arr4d.shape[0]): 
      for j in xrange(arr4d.shape[1]): 
       for k in xrange(arr4d.shape[selected_dim]): 
        if selected_dim==2: 
         arr=arr4d[i,j,k,:] 
        elif selected_dim==3: 
         arr=arr4d[i,j,:,k] 
        do_something(arr) # arr is 1-D and has 10 items 

...但我相信有一些方法來避免嵌套「如果」的一部分，也許還更有效率？就像在循環之前創建此數組的其他視圖，然後迭代這些視圖一樣？處理這種

Answer 1

一種常見的方法是使用np.rollaxis：

def myfun(arr4d, selected_dim): # selected_dim can only be 2 or 3 
    arr4d = np.rollaxis(arr4d, selected_dim) 
    print arr4d.shape # (10, 2, 5, 10) 
    for i in xrange(arr4d.shape[1]): 
     for j in xrange(arr4d.shape[2]): 
      for k in xrange(arr4d.shape[0]): 
       arr=arr4d[k, i, j, :] 
       do_something(arr) # arr is 1-D and has 10 items 

注意np.rollaxis應該返回一個視圖，它實際上並不數組複製。