Numpy slice of arbitrary dimensions
PythonNumpyPython Problem Overview
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?
Python Solutions
Solution 1 - Python
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 - Python
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 - Python
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.