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

Question

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

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

b maps i to j where i and j are the first 2 indexes of a, so ​a[i,j,k]

Desired result after applying b to a is:

[[4 5]
​ [6 7]]

Naive solution:

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

Question: Is there a way to do this using a numpy or scipy routine or routines or fancy indexing?

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). The arrays will be large and this operation will be repeated a large number of times so needs to be scaleable and fast.

What I have tried:

1. Compiling using numba but line profiler suggests my code is slower compared to a similarly sized numpy routine. Also numpy is more widely understood and used.
2. Maintaining pi(a|s) as a sparse matrix (all zero except one 1 per row) b_as_a_matrix and then using einsum but this involves storing and updating the matrix and creates more work (an extra loop over j and sum operation).

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

Answer 1

Numpy arrays can be indexed using other arrays as indices. See also: NumPy selecting specific column index per row by using a list of indexes.

With that in mind, we can vectorize your loop to simply use b for indexing:

>>> 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]])