Search code examples
pythonnumpybooleanlinear-algebrainfinity

How to zero out all indices of an array with np.infs with a boolean array?


I have a matrix with inf values and a boolean array that indicates which values to keep. How do I use the boolean array to zero out all values in the original matrix including the infs but keep all infs corresponding to Trues?

ex

X     = [inf, 1, inf]
        [inf, 2,   4]
        [3,   4,   5]

M     = [1,   0,   0]
        [0,   1,   0]
        [0,   1,   0]

(current output)
M * X = [inf, 0, nan]
        [nan, 2,   0]
        [0,   4,   0]

(desired output)
M * X = [inf, 0,   0]
        [0,   2,   0]
        [0,   4,   0]

Solution

  • Inputs:

    In [77]: X    
    Out[77]: 
    array([[inf,  1., inf],
           [inf,  2.,  4.],
           [ 3.,  4.,  5.]])
    
    In [78]: M 
    Out[78]: 
    array([[1, 0, 0],
           [0, 1, 0],
           [0, 1, 0]])
    

    Approach

    First, we need to invert the mask M and then get the indices using numpy.where; With these indices we can then set the elements in the original array to zero, by indexing into them as follows:

    # inverting the mask
    In [59]: M_not = np.logical_not(M)
    
    In [80]: M_not
    Out[80]: 
    array([[False,  True,  True],
           [ True, False,  True],
           [ True, False,  True]])
    
    # get the indices where `True` exists in array `M_not`
    In [81]: indices = np.where(M_not) 
    
    In [82]: indices 
    Out[82]: (array([0, 0, 1, 1, 2, 2]), array([1, 2, 0, 2, 0, 2]))
    
    # zero out the elements
    In [84]: X[indices] = 0
    
    In [61]: X 
    Out[61]: 
    array([[inf, 0., 0.],
           [0., 2., 0.],
           [0., 4., 0.]])
    

    P.S. inverting the mask should not be understood as matrix inversion. It should be understood as flipping the boolean values (True --> False; False --> True)