Numpy slice of arbitrary dimensions

I would like to slice a numpy array to obtain the i-th index in the last dimension. For a 3D array, this would be:

slice = myarray[:, :, i]

But I am writing a function where I can take an array of arbitrary dimensions, so for a 4D array I'd need myarray[:, :, :, i], and so on. Is there a way I can obtain this slice for any array without explicitly having to write the array dimensions?

22435 次浏览
小开
最佳答案

There is ... or Ellipsis, which does exactly this:

slice = myarray[..., i]

Ellipsis is the python object, if you should want to use it outside the square bracket notation.

小开

Actually, just found the answer. As stated in numpy's documentation this can be done with the slice object. In my particular case, this would do it:

idx = [slice(None)] * (myarray.ndim - 1) + [i]
my_slice = myarray[idx]

The slice(None) is equivalent to choosing all elements in that index, and the last [i] selects a specific index for the last dimension.

小开

In terms of slicing an arbitrary dimension, the previous excellent answers can be extended to:

indx = [slice(None)]*myarray.ndim
indx[slice_dim] = i
sliced = myarray[indx]

This returns the slice from any dimension slice_dim - slice_dim = -1 reproduces the previous answers. For completeness - the first two lines of the above listing can be condensed to:

indx = [slice(None)]*(slice_dim) + [i] + [slice(None)]*(myarray.ndim-slice_dim-1)

though I find the previous version more readable.


提交答案