Indexing and retrieving numpy rows from a multidimensional array

Question

I have a source multidimensional array of shape (a,b,c,c,d) which stores vectors/data of size d , and another array of shape (a,b,e,2) that stores e indices of size 2. 2-dimensional values correspond to the indices 2-3 of the data array (both dimensions of size c ). Note that both arrays share the same a,b dimension sizes.

What I want to do is to use these indices to retrieve rows of size d from the first array. So that, the output array should have size (a,b,e,d) , ie e vectors of size d along the a,b dimensions.

a, b, c, d = 3,5,7,9
e = 11
data = np.random.rand(a,b,c,c,d)
inds = np.random.randint(0,c, size=(a,b,e,2))
    
res = data[:, :, inds[:,:,:,0], inds[:,:,:,1],:]
print(' - Obtained shape:', res.shape) 
print(' - Desired shape:', (a,b,e,d))

# - Obtained shape: (3, 5, 3, 5, 11, 9)
# - Desired shape: (3, 5, 11, 9)

Answer 1

The only way I can think right now is enforcing full fancy indexing by generating range-like indices in all three leading dimensions:

import numpy as np
rng = np.random.default_rng()

a, b, c, d = 3, 5, 7, 9
e = 11
data = rng.uniform(size=(a, b, c, c, d))
inds = rng.integers(0, c, size=(a, b, e, 2))

# generate open index meshes to reduce memory for at least here
aind, bind, _, = np.ogrid[:a, :b, :e]
res = data[aind, bind, inds[..., 0], inds[..., 1], :]

print(' - Obtained shape:', res.shape) 
print(' - Desired shape:', (a, b, e, d))

Random check to see that the values are correct too:

sample_index_pos = (1, 1, 8)  # <-> (a, b, e)
c_inds = inds[sample_index_pos] # <-> (c, c)
expected = data[sample_index_pos[:2] + tuple(c_inds)]
have = res[sample_index_pos]
print(np.array_equal(expected, have))
# True