How can I update a numpy array with index in another numpy array

Question

I have

• an numpy.array a of shape (n1, n2, n3, n4)
• an index array idx of shape (n1, n2, i1)

what I want to do is like the code below

for i in range(n1):
    for j in range(n2):
        for k in range(i1):
            b[i, j, k, :] = a[i, j, idx[i, j, k], :]

if there is a numpy function to achieve this without for loop?

Answer 1

Using as starting point:

import numpy as np

n1, n2, n3, n4, i1 = range(2, 7)

a = np.random.randint(10, size=(n1, n2, n3, n4))
idx = np.random.randint(n3, size=(n1, n2, i1))
b = np.zeros_like(a, shape=(n1, n2, i1, n4))

In general you can do the following:

I, J, K = np.ogrid[:n1, :n2, :i1]
b[I, J, K] = a[I, J, idx]

Here the I J and K arrays are the equivalent of the loop variables i j and k. Their shapes have to be in agreement with the shape of idx.

In case b has shape (n1, n2, i1, n4) then you might as well do:

I, J, _ = np.ogrid[:n1, :n2, :1]
b = a[I, J, idx]

Or alternatively without ogrid:

b = np.take_along_axis(a, idx[...,np.newaxis], axis=2)

Here newaxis is used to insert a length-1 axis to allow broadcasting. Check out the numpy indexing docs for more info.