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