NumPy:在矢量化赋值期间Evaulate索引数组

Question

我想对这个NumPy操作进行矢量化:

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

其中idx包含x切片的axis-0索引.有一些简单的方法可以做到这一点吗？

Answer 1

您可以使用:

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

这是测试:

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)