由INT選擇numpy數組軸

Question

我想系統地訪問一個numpy數組的軸。例如，假設我有一個數組

a = np.random.random((10, 10, 10, 10, 10, 10, 10)) 
# choosing 7:9 from axis 2 
b = a[:, :, 7:9, ...] 
# choosing 7:9 from axis 3 
c = a[:, :, :, 7:9, ...] 

如果我有高維數組，則鍵入冒號會變得非常重複。現在，我想一些功能choose_from_axis這樣

# choosing 7:9 from axis 2 
b = choose_from_axis(a, 2, 7, 9) 
# choosing 7:9 from axis 3 
c = choose_from_axis(a, 3, 7, 9) 

所以，基本上，我想用一個號碼訪問軸。我知道如何做到這一點的唯一方法是來回使用rollaxis，但我正在尋找一種更直接的方式來做到這一點。

Answer 1

聽起來像你可能正在尋找take：

>>> a = np.random.randint(0,100, (3,4,5)) 
>>> a[:,1:3,:] 
array([[[61, 4, 89, 24, 86], 
     [48, 75, 4, 27, 65]], 

     [[57, 55, 55, 6, 95], 
     [19, 16, 4, 61, 42]], 

     [[24, 89, 41, 74, 85], 
     [27, 84, 23, 70, 29]]]) 
>>> a.take(np.arange(1,3), axis=1) 
array([[[61, 4, 89, 24, 86], 
     [48, 75, 4, 27, 65]], 

     [[57, 55, 55, 6, 95], 
     [19, 16, 4, 61, 42]], 

     [[24, 89, 41, 74, 85], 
     [27, 84, 23, 70, 29]]]) 

這也會支持元組索引。例如：

>>> a = np.arange(2*3*4).reshape(2,3,4) 
>>> a 
array([[[ 0, 1, 2, 3], 
     [ 4, 5, 6, 7], 
     [ 8, 9, 10, 11]], 

     [[12, 13, 14, 15], 
     [16, 17, 18, 19], 
     [20, 21, 22, 23]]]) 
>>> a[:,:,(0,1,3)] 
array([[[ 0, 1, 3], 
     [ 4, 5, 7], 
     [ 8, 9, 11]], 

     [[12, 13, 15], 
     [16, 17, 19], 
     [20, 21, 23]]]) 
>>> a.take((0,1,3), axis=2) 
array([[[ 0, 1, 3], 
     [ 4, 5, 7], 
     [ 8, 9, 11]], 

     [[12, 13, 15], 
     [16, 17, 19], 
     [20, 21, 23]]]) 

Answer 2

您可以構建一個切片對象，沒有工作：

def choose_from_axis(a, axis, start, stop): 
    s = [slice(None) for i in range(a.ndim)] 
    s[axis] = slice(start, stop) 
    return a[s] 

例如，以下均可以得到相同的結果：

x[:,1:2,:] 
choose_from_axis(x, 1, 1, 2) 

# [[[ 3 4 5]] 
# [[12 13 14]] 
# [[21 22 23]]] 

一樣的問題的例子：

a = np.random.random((10, 10, 10, 10, 10, 10, 10)) 
a0 = a[:, :, 7:9, ...] 
a1 = choose_from_axis(a, 2, 7, 9) 

print np.all(a0==a1) # True