Multi-dimensional indexing in Numpy

Question

I checked the Numpy docs and so on this but I couldn't find an answer. Maybe it can't be done.

Basically I have an array probabilities of shape (3,4,5). In the 3rd dimension there are 5 elements, which together sum to 1. The element index in the third dimension that I want corresponds to the values in the array index of shape (3,4). Makes sense?

So if probabilities[0,0,:] is equal to [0.1, 0.1, 0.2, 0.4, 0.2] and index[0,0] is equal to 2 , then I want the 3rd element which is 0.2 .

I tried probabilities[index] and other things, but no luck.

Can this be done without a loop?

Answer 1

Make a sample array:

In [291]: A = np.arange(2*3*4).reshape(2,3,4)
In [292]: A[0,0,:]
Out[292]: array([0, 1, 2, 3])
In [293]: A[0,0,2]
Out[293]: 2

Make a sample idx:

In [294]: idx = np.random.randint(0,4,(2,3),int)
In [295]: idx
Out[295]: 
array([[0, 3, 0],
       [1, 0, 1]])

These are index values for the 3rd dimension. Make arrays for indexing on the 1st 2 dimensions:

In [299]: I,J=np.ix_(np.arange(A.shape[0]),np.arange(A.shape[1]))
In [300]: I,J
Out[300]: 
(array([[0],
        [1]]), array([[0, 1, 2]]))
In [301]: A[I,J,idx]
Out[301]: 
array([[ 0,  7,  8],
       [13, 16, 21]])

test:

In [302]: A[0,1,3]
Out[302]: 7
In [304]: A[1,2,1]
Out[304]: 21

There are various ways of getting those I,J . np.ix_ is an easy one. So is np.ogrid , np.mgrid or even np.meshgrid .

In [306]: I,J = np.mgrid[0:2,0:3]
In [307]: I,J
Out[307]: 
(array([[0, 0, 0],
        [1, 1, 1]]), array([[0, 1, 2],
        [0, 1, 2]]))
In [308]: A[I,J,idx]
Out[308]: 
array([[ 0,  7,  8],
       [13, 16, 21]])