CODEWITHSUNDEEP
pythonnumpyindexingslicedimensions
EliteKaffee
EliteKaffee

Reputation: 109

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?

Upvotes: 1

Views: 209

Answers (1)

sshashank124
sshashank124

Reputation: 32197

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

Upvotes: 1

Related Questions