使用索引之间的一维查找对 3d numpy 数组进行切片

Question

import numpy as np
a = np.arange(12).reshape(2, 3, 2)
b = np.array([2, 0])

b 将 i 映射到 j，其中 i 和 j 是 a 的前两个索引，所以 a[i,j,k]

将 b 应用于 a 后的预期结果是：

[[4 5]
 [6 7]]

天真的解决方案：

c = np.empty(shape=(2, 2), dtype=int)
for i in range(2):
   j = b[i]
   c[i, :] = a[i, j, :]

问题：有没有办法使用 numpy 或 scipy 例程或例程或花哨的索引来做到这一点？

应用：强化学习有限 MDP，其中 b 是确定性策略向量 pi(a|s)，a 是状态转移概率 p(s'|s,a)，c 是该策略的状态转移矩阵向量 p(s'|s)。数组会很大，这个操作会重复很多次，所以需要可缩放和快速。

我尝试过的：

使用 numba 进行编译，但行分析器表明我的代码与类似大小的 numpy 例程相比速度较慢。 numpy 也得到了更广泛的理解和使用。
将 pi(a|s) 维护为稀疏矩阵（除每行一个 1 外全为零）b_as_a_matrix 然后使用 einsum 但这涉及存储和更新矩阵并创建更多工作（一个额外的循环j 和求和运算）。

c = np.einsum('ij,ijk->ik', b_as_a_matrix, a)

Answer 1

Numpy 数组可以使用其他数组作为索引进行索引。另见：NumPy selecting specific column index per row by using a list of indexes.

考虑到这一点，我们可以将您的循环向量化以简单地使用 b 进行索引：

>>> import numpy as np
>>> a = np.arange(12).reshape(2, 3, 2)
>>> b = np.array([2, 0])
>>> i = np.arange(len(b))
>>> i
array([0, 1])
>>> a[i, b, :]
array([[4, 5],
       [6, 7]])

使用索引之间的一维查找对 3d numpy 数组进行切片

Slice a 3d numpy array using a 1d lookup between indices

python

numpy

scipy

reinforcement-learning

numba