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.
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.