Search code examples
pythonarrayspython-3.xnumpycompare

How to compare two numpy arrays with some NaN values?


I need to compare some numpy arrays which should have the same elements in the same order, excepting for some NaN values in the second one.

I need a function more or less like this:

def func( array1, array2 ):
    if ???:
        return True
    else:
        return False

Example:

x = np.array( [ 1, 2, 3, 4, 5 ] )
y = np.array( [ 11, 2, 3, 4, 5 ] )
z = np.array( [ 1, 2, np.nan, 4, 5] )

func( x, z ) # returns True
func( y, z ) # returns False

The arrays have always the same length and the NaN values are always in the third one (x and y have always numbers only). I can imagine there is a function or something already, but I just don't find it.

Any ideas?


Solution

  • You can use masked arrays, which have the behaviour you're asking for when combined with np.all:

    zm = np.ma.masked_where(np.isnan(z), z)
    
    np.all(x == zm) # returns True
    np.all(y == zm) # returns False
    

    Or you could just write out your logic explicitly, noting that numpy has to use | instead of or, and the difference in operator precedence that results:

    def func(a, b):
        return np.all((a == b) | np.isnan(a) | np.isnan(b))