CODEWITHSUNDEEP
pythonnumpy
Nir
Nir

Reputation: 1894

Numpy slice of first arbitrary dimensions

There is this great Question/Answer about slicing the last dimension: Numpy slice of arbitrary dimensions: for slicing a numpy array to obtain the i-th index in the last dimension, one can use ... or Ellipsis,

slice = myarray[...,i]

What if the first N dimensions are needed ?

For 3D myarray, N=2:

slice = myarray[:,:,0]

For 4D myarray, N=2:

slice = myarray[:,:,0,0]

Does this can be generalized to an arbitrary dimension?

Upvotes: 0

Views: 842

Answers (1)

Hans Musgrave
Hans Musgrave

Reputation: 7131

I don't think there's any built-in syntactic sugar for that, but slices are just objects like anything else. The slice(None) object is what is created from :, and otherwise just picking the index 0 works fine.

myarray[(slice(None),)*N+(0,)*(myarray.ndim-N)]

Note the comma in (slice(None),). Python doesn't create tuples from parentheses by default unless the parentheses are empty. The comma signifies that don't just want to compute whatever's on the inside.

Slices are nice because they give you a view into the object instead of a copy of the object. You can use the same idea to, e.g., iterate over everything except the N-th dimension on the N-th dimension. There have been some stackoverflow questions about that, and they've almost unanimously resorted to rolling the indices and other things that I think are hard to reason about in high-dimensional spaces. Slice tuples are your friend.

From the comments, @PaulPanzer points out another technique that I rather like.

myarray.T[(myarray.ndim-N)*(0,)].T

First, transposes in numpy are view-operations instead of copy-operations. This isn't inefficient in the slightest. Here's how it works:

  1. Start with myarray with dimensions (0,...,k)
  2. The transpose myarray.T reorders those to (k,...,0)
  3. The whole goal is to fix the last myarray.ndim-N dimensions from the original array, so we select those with [(myarray.ndim-N)*(0,)], which grabs the first myarray.ndim-N dimensions from this array.
  4. They're in the wrong order. We have dimensions (N-1,...,0). Use another transpose with .T to get the ordering (0,...,N-1) instead.

Upvotes: 4

Related Questions