在二维numpy数组中使用三维样式切片

Question

我有一个以numpy数组（a）作为输入的函数。此阵列可以是二维或三维阵列，具体取决于数学计算。有一个整数m可以是任何数字，除非数组是2D，否则m的值总是0。我想传递一个silce到另一个函数。因为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中实现它？我使用python2。你知道吗

Answer 1

使用^{}似乎是一个很好的设置，因为我们可以强制它成为3D，然后简单地沿着最后一个轴切片第m个索引，就像这样-

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

它仍然是阵列的一个视图，因此没有效率损失！你知道吗

案例运行

1）二维：

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