Search code examples
pythonnumpyindexingupdatesmask

Update by an index and a mask?


I have a big 2-dimensional array which I access by indexes. I want to update only the values of the indexed array which are not zero.

arrayx = np.random.random((10,10))

Let's say I have indexes (this is just example, the actual indexes are generated by separate process):

idxs = np.array([[4],
        [5],
        [6],
        [7],
        [8]]), np.array([[5, 9]])

Given these indexes, this should work, but it doesn't.

arrayx[idxs]

array([[0.7 , 0.1 ],
       [0.79, 0.51],
       [0.  , 0.8 ],
       [0.82, 0.32],
       [0.82, 0.89]], dtype=float16)

// note from editor: '<>' is equivalent to '!='
// but I agree that '>' 0 is more correct
// mask = mapx[idxs] <> 0 // original
mask = arrayx[idxs] > 0 // better

array([[ True,  True],
       [ True,  True],
       [False,  True],
       [ True,  True],
       [ True,  True]])

arrayx[idxs][mask] += 1

However, this does not update the array. How can I solve this?


Solution

  • A simple one with np.where with a mask as the first input to choose and assign -

    mapx[idxs] = np.where(mask,mapx[idxs]+1,mapx[idxs])
    

    Custom update values

    The second argument (here mapx[idxs]+1) could be edited to any complex update that you might be doing for the masked places corresponding to True ones in mask. So, let's say you were doing an update for the masked places with :

    mapx[idxs] += x * (A - mapx[idxs])
    

    Then, replace the second arg to mapx[idxs] + x * (A - mapx[idxs]).


    Another way would be to extract integer indices off True ones in mask and then create new idxs that is selective based on the mask, like so -

    r,c = np.nonzero(mask)
    idxs_new = (idxs[0][:,0][r], idxs[1][0][c])
    mapx[idxs_new] += 1
    

    The last step could be edited similarly for a custom update. Just use idxs_new in place of idxs to update.