Search code examples
pythonnumpyindexingslicedimensions

Pythonic way to index without losing dimensions


data = np.zeros((5, 16, 16))

I would like to access a slice of size 1 of this numpy data array. It might be helpful to define what I don't want:

>>> data[2].shape
(16, 16)

Instead, I want to keep the dimensionality the same. A couple of ways come to mind here:

>>> np.expand_dims(data[2], 0).shape
(1, 16, 16)

>>> data[2:3].shape
(1, 16, 16)

>>> data[None, 2].shape
(1, 16, 16)

Which one of these options is more pythonic, or is there a better solution here?


Solution

  • You can also do it with a list of indices with a single element:

    >>> data[[2]].shape
    (1, 16, 16)
    

    As for which is more pythonic, that is more opinion-based than anything.


    Note: This method will create a copy of the data instead of a view into the same data since arbitrary indices might not result in a contiguous view. This is explained in detail in the question Avoid copying when indexing a numpy arrays using lists