在 2D 数字数组中使用 3D 样式切片

Question

我有一个函数，它接受一个 numpy 数组 (A) 作为输入。此数组可以是 2d 或 3d 数组，具体取决于数学计算。有一个整数 m 可以是任何数字，除了当数组是 2D 时，m 的值将始终为 0。我想将 A 的 sill 传递给另一个函数。由于 A 既可以是 3D 也可以是 2D，所以我尝试了 3D 样式切片。

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

这给了我一个IndexError当 A 是 2D 时

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