Numpy slice of arbitrary dimensions

Question

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?

Solution 1

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.

Solution 2

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.

Solution 3

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.

Comments

[Niklas B.]

I was about to suggest this, but actually Sebastian's approach is much more elegant! Also, I don't think that your code actually works, you are multiplying a list with a tuple and then you are concatenating a bare index without a slice() object. Maybe you meant [slice(None)] * myarray.shape[-1] + [slice(None,i)]

[tiago]

Thank you, I didn't know about this. My own answer allows for more generic slices, but yours is really cool.

[seberg]

@tiago Indeed you can always build a slicing tuple. >ou can use Ellipsis in there too...

[seberg]

@NiklasB. The last [i] is correct, but the first part must be (myarray.ndim - 1) for the multiplier if Ellipsis is not used.

[tiago]

@Sebastian, you're right, I tested using a number and then forgot to put len(myarray.shape[:-1]). I just fixed it with your suggestion.

[Kyle Heuton]

I love this answer, but I don't like that the slice object is immediately over-written

[TheLoneDeranger]

I really appreciate this, because pscript doesn't support ellipses yet (strangely, since JS does...). So yeah, I really appreciate ya'll taking the time to make and improve this alternative anser ( <-- sorry don't have that key, too much fps :( :( )

[dashesy]

I find idx=[np.s_[:]] * data.ndim to be more intuitive, I can then change one of those to something else, like idx[-1]=np.s_[:2] to slice the last axis differently.

[user4779]

Couldn't find this answer anywhere and surprised it's not in the official python sections on multidimensional arrays. Thank you very much!

[F. Remonato]

Very good solution, just know that as of today on Numpy 1.18.1 the [Ellipsis] part needs to be changed to [slice(None)], otherwise a "IndexError: an index can only have a single ellipsis ('...')" is raised.

[Moot]

This answer is great, except that it assigns the sliced array to the builtin slice - I'd suggest to rename that variable to avoid issues.