Search code examples
pythonarraysnumpyconditional-statements

Replacing Numpy elements if condition is met


I have a large numpy array that I need to manipulate so that each element is changed to either a 1 or 0 if a condition is met (will be used as a pixel mask later). There are about 8 million elements in the array and my current method takes too long for the reduction pipeline:

for (y,x), value in numpy.ndenumerate(mask_data): 

    if mask_data[y,x]<3: #Good Pixel
        mask_data[y,x]=1
    elif mask_data[y,x]>3: #Bad Pixel
        mask_data[y,x]=0

Is there a numpy function that would speed this up?


Solution

  • >>> import numpy as np
    >>> a = np.random.randint(0, 5, size=(5, 4))
    >>> a
    array([[4, 2, 1, 1],
           [3, 0, 1, 2],
           [2, 0, 1, 1],
           [4, 0, 2, 3],
           [0, 0, 0, 2]])
    >>> b = a < 3
    >>> b
    array([[False,  True,  True,  True],
           [False,  True,  True,  True],
           [ True,  True,  True,  True],
           [False,  True,  True, False],
           [ True,  True,  True,  True]], dtype=bool)
    >>> 
    >>> c = b.astype(int)
    >>> c
    array([[0, 1, 1, 1],
           [0, 1, 1, 1],
           [1, 1, 1, 1],
           [0, 1, 1, 0],
           [1, 1, 1, 1]])
    

    You can shorten this with:

    >>> c = (a < 3).astype(int)