numpy 的 assert_array_equal 是不是验证复杂的 NaN 完全相等？

Posted 2023-02-22

【中文标题】numpy 的 assert_array_equal 是不是验证复杂的 NaN 完全相等？

【英文标题】：Does numpy's assert_array_equal verify that complex NaNs are exactly equal?numpy 的 assert_array_equal 是否验证复杂的 NaN 完全相等？

【发布时间】：2020-05-05 14:50:47

【问题描述】：

我想要两个包含NaN 的复数的numpy数组是否完全相等。

具体来说，我想验证对于NaN，不仅两个数组中都有NaN，而且NaN 的实部和虚部在值上匹配。

有人知道assert_array_equal 是这样做还是我必须自己检查一下？

【问题讨论】：

如果你提供一个例子会有所帮助。

这似乎很容易在提示下进行测试。对我来说，np.testing.assert_array_equal(np.array([complex(2.3, np.nan)]), np.array([complex(np.nan, 2.3)])) 不会提高。

@MarkDickinson 当然！我自己应该想到的……非常感谢。

【参考方案1】：

Numpy 只是检查是否有 NaN 值 in the same positions 并为此目的使用 np.isnan。这里实部或虚部是否包含 NaN 值并不重要：

>>> np.isnan(np.sqrt(-1.) + 1j)
True
>>> np.isnan(np.sqrt(-1.) * 1j)
True

同样适用于两个数组a 和b：

>>> a = np.zeros(3, dtype=np.complex128)
>>> b = a.copy()
>>> a[0] = np.sqrt(-1.) + 1j
>>> b[0] = np.sqrt(-1.) * 1j
>>> a
array([nan+1.j,  0.+0.j,  0.+0.j])
>>> b
array([nan+nanj,  0. +0.j,  0. +0.j])
>>> np.testing.assert_array_equal(a, b) is None
True

【参考方案2】：

来自comparing numpy arrays containing NaN

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

要比较复数，只需检查实部和虚部是否相等。例如：

a = np.array([1+2j, 3+4j, np.nan+6j])
b = np.array([1+2j, 3+4j, np.nan+5j])

eq(a.real, b.real) & eq(a.imag, b.imag)
False

编辑：或者你可以使用np.allclose(a.real, b.real, equal_nan=True) & np.allclose(a.imag, b.imag, equal_nan=True)。