Проверяет ли assert_array_equal numpy, что комплексные NaN точно равны?

Question

Я хочу, чтобы два numpy массива с комплексными числами, содержащими NaN, были в точности равны.

В частности, я хочу убедиться, что для NaN не только есть NaN в обоих массивы, но также и то, что действительная и мнимая части NaN совпадают по значению.

Кто-нибудь знает, делает ли assert_array_equal это или я должен сам это проверить?

Answer 1

С сравнение numpy массивов, содержащих 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).

Answer 2

Numpy просто проверяет, есть ли NaN значения в тех же позициях и использует 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