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

Question

I have a function that takes a numpy array (A) as input. This array could be a 2d or a 3d array depending on a mathematical calculation. There is an integer m which could be any number, except when the array is 2D, the value of m will always be 0. I want to pass a silce of A to another function. Since A can be both 3D or 2D, I tried 3D style slicing.

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

This gives me an IndexError when A is 2D

IndexError: too many indices for array

I want to pass the full 2D array to fun2 if A is 2D, like it happens in MATLAB. How can it be done in Python? I use Python 2.

Answer 1

Seems like a good setup to use np.atleast_3d as we can force it to be 3D and then simply slice the m-th index along the last axis, like so -

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

It's still a view into the array, so no efficiency lost there!

Case runs

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