【Python】Numpy の assert_array_equal が思ったように動かない原因は dtype にある

概要

マニュアルによると numpy.testing.assert_array_equal は NaN (np.nan) 同士を同じものとして比較してくれる、はずだが動かなかった。

原因は、dtype が object のとき (mixed な ndarray のとき) は思ったように動かない、と判明

環境

Linux 系 OS (詳細未確認) Conda 版 Python v3.7 および Numpy v1.8.1

$ python
Python 3.7.0 (default, Oct  9 2018, 10:31:47) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.

動くケース

通常の比較では、NaN 同士の比較は False になるが、assert_array_equal は同じものとして比較してくれる。
ただし、float の配列なら。

>>> import numpy as np
>>> from numpy.testing import assert_array_equal as aae

>>> a = np.array([1, np.nan])
>>> b = np.array([1, np.nan])

>>> a.dtype
dtype('float64')
>>> b.dtype
dtype('float64')

>>> a == b
array([ True, False])
>>> aae(a, b)
>>> 

動かないケース

a および b を作るときに dtype を指定しているのは、そうしないと文字列の配列になってしまうからである。
object 型として作られたこれらの配列に対して、assert_array_equal は今度は fail する。

>>> a = np.array([1,'test',np.nan], dtype=object)
>>> b = np.array([1,'test',np.nan], dtype=object)

>>> a.dtype
dtype('O')
>>> b.dtype
dtype('O')

>>> a == b
array([ True,  True, False])
>>> aae(a,b)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../anaconda3/envs/opt/lib/python3.7/site-packages/numpy/testing/_private/utils.py", line 936, in assert_array_equal
    verbose=verbose, header='Arrays are not equal')
  File ".../anaconda3/envs/opt/lib/python3.7/site-packages/numpy/testing/_private/utils.py", line 846, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Arrays are not equal

Mismatched elements: 1 / 3 (33.3%)
 x: array([1, 'test', nan], dtype=object)
 y: array([1, 'test', nan], dtype=object)
>>>