python - NumPy:在向量化分配期间评估索引数组

Question

我想向量化这个 NumPy 操作:

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

其中 idx 包含 x 切片的轴 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)