python - Comparing NumPy arrays so that NaNs compare equal

Question

Is there an idiomatic way to compare two NumPy arrays that would treat NaNs as being equal to each other (but not equal to anything other than a NaN).

For example, I want the following two arrays to compare equal:

np.array([1.0, np.NAN, 2.0])
np.array([1.0, np.NAN, 2.0])

and the following two arrays to compare unequal:

np.array([1.0, np.NAN, 2.0])
np.array([1.0, 0.0, 2.0])

I am looking for a method that would produce a scalar Boolean outcome.

The following would do it:

np.all((a == b) | (np.isnan(a) & np.isnan(b)))

but it's clunky and creates all those intermediate arrays.

Is there a way that's easier on the eye and makes better use of memory?

P.S. If it helps, the arrays are known to have the same shape and dtype.

Answer 1

If you really care about memory use (e.g. have very large arrays), then you should use numexpr and the following expression will work for you:

np.all(numexpr.evaluate('(a==b)|((a!=a)&(b!=b))'))

I've tested it on very big arrays with length of 3e8, and the code has the same performance on my machine as

np.all(a==b)

and uses the same amount of memory