pythonnumpy

Broadcast advanced indexing numpy


I have an array of values for example

x = array([[[-0.78867513, -0.21132487,  0.        ,  0.78867513,  0.21132487,    0.        ,  0.        ,  0.        ,  0.        ],
            [ 0.        , -0.78867513, -0.21132487,  0.        ,  0.78867513,    0.21132487,  0.        ,  0.        ,  0.        ],
            [ 0.        ,  0.        ,  0.        , -0.78867513, -0.21132487,    0.        ,  0.78867513,  0.21132487,  0.        ],
            [ 0.        ,  0.        ,  0.        ,  0.        , -0.78867513,   -0.21132487,  0.        ,  0.78867513,  0.21132487]],
           [[-0.78867513, -0.21132487,  0.        ,  0.78867513,  0.21132487,    0.        ,  0.        ,  0.        ,  0.        ],
            [ 0.        , -0.78867513, -0.21132487,  0.        ,  0.78867513,    0.21132487,  0.        ,  0.        ,  0.        ],
            [ 0.        ,  0.        ,  0.        , -0.78867513, -0.21132487,    0.        ,  0.78867513,  0.21132487,  0.        ],
            [ 0.        ,  0.        ,  0.        ,  0.        , -0.78867513,   -0.21132487,  0.        ,  0.78867513,  0.21132487]]])

I want in use advanced indexing to pull out the nonzero values. I know the indices of the nonzero values so

idx = array([[4, 3, 1, 0],
             [5, 4, 2, 1],
             [7, 6, 4, 3],
             [8, 7, 5, 4]])

The desired result would be something like

x[idx] = array([[[-0.78867513, -0.21132487,  0.78867513,  0.21132487],
                 [-0.78867513, -0.21132487,  0.78867513,  0.21132487],
                 [-0.78867513, -0.21132487,  0.78867513,  0.21132487],
                 [-0.78867513, -0.21132487,  0.78867513,  0.21132487]],
                [[-0.78867513, -0.21132487,  0.78867513,  0.21132487],
                 [-0.78867513, -0.21132487,  0.78867513,  0.21132487],
                 [-0.78867513, -0.21132487,  0.78867513,  0.21132487],
                 [-0.78867513, -0.21132487,  0.78867513,  0.21132487]]])

The actual x array is much larger along the first dimension, but the nonzero structure is always indicated by idx so I need it to broadcast along the first dimension. Is this possible?

EDIT: To be clear x along the first dimension contains a nested list of 4 x 9 array. idx then has the nonzero entries row-for-row. Notice in the first row of the both 4 x 9 nested arrays in x that the 4 3 1 0 entries are nonzero.


Solution

  • Try this one:

    x[:,np.arange(idx.shape[0])[:,None],idx]
    

    Using this technique every element in np.arange(idx.shape[0])[:,None] (which has shape (idx.shape[0], 1) and therefore is a column vector) will be broadcast with every row in idx. This will then be used for all entries along x's first axis.