Numpy common slicing for different dimensions-CodePudding

I have a function performs an operation on 2D arrays. I would like to expand its functionality to be able to work with 3D arrays, for which I would like to know if there is a way to access the vectors of each 2D array regardless of whether the input is a 3D array or a 2D array.

For example, if I have the following 2D matrix:

>>>arr2d=array([[0, 0],
                [1, 1]])

I can access the last vector using:

>>>arr2d[-1]
array([1, 1])

And if I have a 3D array like this:

>>>arr3d=array([[[ 0,  0],
                 [ 1,  1]],

                [[ 3,  0],
                 [ 2,  7]],

                [[ 9,  5],
                 [ 8,  6]],

                [[20,  4],
                 [ 6,  5]]])

I can access the last vector of each 2D submatrix using:

>>>arr3d[:,-1]
array([[1, 1],
       [2, 7],
       [8, 6],
       [6, 5]])

I would like to know if there is a common slice that I can use on both arrays to get the above results, i.e. something like the following, with some_slice being the same in each case:

>>>arr2d[some_slice]
array([1, 1])

>>>arr3d[some_slice]
array([[1, 1],
       [2, 7],
       [8, 6],
       [6, 5]])

CodePudding user response：

Use Ellipsis as below:

print(arr2d[..., -1, :])
print(arr3d[..., -1, :])

Output

[1 1]
[[1 1]
 [2 7]
 [8 6]
 [6 5]]

From the documentation (emphasis mine):

Ellipsis expands to the number of : objects needed for the selection tuple to index all dimensions. In most cases, this means that the length of the expanded selection tuple is x.ndim. There may only be a single ellipsis present

CodePudding user response：

you can use the ... slice in this way:

arr2d[...,-1,:]
arr3d[...,-1,:]

in both cases you are slicing the last-but-one axis. In general you can set custom_slice = (Ellipsis, -1, slice(None)) and apply it to bot arrays:

arr2d[custom_slice]
arr3d[custom_slice]