Indexing a multi-dimensional array partially

Question

import numpy as np
arr = np.random.rand(50,3,3,3,16)
ids = (0,0,2,10)
b = arr[:, ids]  # don't work
b = arr[:, *ids]  # don't work
b = arr[:][ids]  # don't work
b = arr[:, tuple(ids)]  # don't work
b = arr[: + ids]  # don't work, obviously..
# b = arr[:,0,0,2,10].shape  # works (desired outcome)

I know there have been several questions about this, like Tuple as index of multidimensional array or Unpacking tuples/arrays/lists as indices for Numpy Arrays but none of them work for my case. Basically I want to index everything in the first axis, of specified "columns" in the rest of the axes (see the last line of my code). The desired output shape should be (50,) in this case.

But I want to index with a tuple/list of ids because I need to iterate through them, for example:

all_ids = ((0,0,0,2), (0,0,0,6), (1,1,0,2), (1,1,0,6),
           (2,2,0,2), (2,2,0,6), (2,2,2,2), (2,2,2,6))
c = 0
for id in all_ids:
    c += arr[:, id].sum()

Psidom Psidom · Accepted Answer · 2018-11-07T18:02:32

Add slice(None) to first dimension in ids and then subset:

arr[(slice(None),) + ids].shape
# (50,)

where:

(slice(None),) + ids
# (slice(None, None, None), 0, 0, 2, 10)

Notice slice(None, None, None) is equivalent to :, i.e. slice all. You can read docs on using slice object for indexing here.

Indexing a multi-dimensional array partially

1 Answers