在2D numpy數組中使用3D樣式切片

Question

我有一個函數，需要一個numpy數組（A）作為輸入。 根據數學計算，此數組可以是2d或3d數組。 有一個整數m，該整數可以是任意數字，但當數組為2D時，m的值將始終為0。我想將A的silc傳遞給另一個函數。 由於A可以是3D或2D，因此我嘗試了3D樣式切片。

def fun(A):
    ... some code
    ans = fun2(A[:,:,m]) #The value of m is 0 if A is 2D

當A為2D時，這給我一個IndexError

IndexError: too many indices for array

如果A是2D，我想將完整的2D數組傳遞給fun2，就像在MATLAB中一樣。 如何在Python中完成？ 我使用Python 2。

Answer 1

似乎是使用np.atleast_3d的好設置，因為我們可以將其強制為3D ，然后簡單地沿最后一個軸切片第m個索引，如下所示-

np.atleast_3d(A)[...,m] # Or np.atleast_3d(A)[:,:,m]

它仍然是陣列的視圖，因此不會損失任何效率！

案例運行

1）2D：

In [160]: A = np.random.randint(11,99,(4,5))

In [161]: np.atleast_3d(A)[...,0]
Out[161]: 
array([[13, 84, 38, 15, 26],
       [64, 91, 29, 11, 48],
       [25, 66, 77, 14, 87],
       [59, 96, 98, 30, 88]])

In [162]: A
Out[162]: 
array([[13, 84, 38, 15, 26],
       [64, 91, 29, 11, 48],
       [25, 66, 77, 14, 87],
       [59, 96, 98, 30, 88]])

2）3D：

In [163]: A = np.random.randint(11,99,(4,3,5))

In [164]: np.atleast_3d(A)[...,1]
Out[164]: 
array([[34, 81, 66],
       [56, 20, 25],
       [45, 36, 64],
       [82, 64, 31]])

In [165]: A[:,:,1]
Out[165]: 
array([[34, 81, 66],
       [56, 20, 25],
       [45, 36, 64],
       [82, 64, 31]])