Broadcast advanced indexing numpy

Question

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.

Answer 1

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.

Answer 2

I tried this one liner for your problem and it seems to do the job without needing idx. You may need to change the parameter in .reshape() according to the size of your problem.

np.array(filter(lambda x: x!=0, x.ravel())).reshape(-1, 4, 4)

It flattens the array, removes the zeroes and then changes it back to the required shape.

Here's another version which is probably more efficient as it does not use the filter function and uses boolean indexing for numpy arrays instead

y = x.ravel()
z = y[y!=0].reshape(-1, 4, 4)

EDIT:

While playing around with numpy I discovered yet another way to do it.

x[x!=0].reshape(-1, 4, 4)

And here's the performance of all three method:

• Method 1: 10000 loops, best of 3: 21.2 µs per loop
• Method 2: 100000 loops, best of 3: 2.42 µs per loop
• Method 3: 100000 loops, best of 3: 1.97 µs per loop