使用索引之間的一維查找對 3d numpy 數組進行切片

Question

使用索引之間的一維查找對 3d numpy 數組進行切片

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

b 將 i 映射到 j，其中 i 和 j 是 a 的前 2 個索引，所以 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 例程或例程或花哨的索引來做到這一點？

Application: Reinforcement Learning finite MDPs where b is a deterministic policy vector pi(a|s), a is the state transition probabilities p(s'|s,a) and c is the state transition matrix for that policy vector p(s' |s)。 arrays 將很大，並且此操作將重復很多次，因此需要可擴展且快速。

我試過的：

1. 使用 numba 編譯，但 line profiler 表明我的代碼與類似大小的 numpy 例程相比要慢。 numpy 也得到了更廣泛的理解和使用。
2. 將 pi(a|s) 保持為稀疏矩陣（除每行一個 1 外全為零） b_as_a_matrix 然后使用 einsum 但這涉及存儲和更新矩陣並創建更多工作（對 j 和求和運算的額外循環）。

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

Answer 1

Numpy arrays 可以使用其他 arrays 作為索引進行索引。 另請參閱： NumPy 通過使用索引列表選擇每行的特定列索引。

考慮到這一點，我們可以將您的循環矢量化以簡單地使用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]])