NumPy: Evaulate index array during vectorized assignment

Question

I would like to vectorize this NumPy operation:

for j in range(yt):
    for i in range(xt):
        y[j, i] = x[idx[j, i], j, i]

where idx contains axis-0 index to an x slice. Is there some simple way to do this?

Answer 1

You can use:

J, I = np.ogrid[:yt, :xt]
x[idx, J, I]

Here is the test:

import numpy as np

yt, xt = 3, 5
x = np.random.rand(10, 6, 7)
y = np.zeros((yt, xt))
idx = np.random.randint(0, 10, (yt, xt))

for j in range(yt):
    for i in range(xt):
        y[j, i] = x[idx[j, i], j, i]

J, I = np.ogrid[:yt, :xt]
np.all(x[idx, J, I] == y)