Очистите массив np от NaN, удалив записи в другом массиве соответственно

Question

У меня есть два массива numpy, один из которых содержит около 1% NaN.

a = np.array([-2,5,nan,6])
b = np.array([2,3,1,0])

Я бы хотел вычислить среднеквадратическую ошибку a и b, используя sklearn 'mean_squared_error.

Итак, мой вопрос: каков питонный способ удаления всех NaN из a, одновременно удаляя все соответствующие записи из b настолько эффективно, насколько это возможно?

Answer 1

Вы можете просто использовать ванильный NumPy's np.nanmean для этой цели:

In [136]: np.nanmean((a-b)**2)
Out[136]: 18.666666666666668

Если этого не существует или вы действительно хотите использовать метод sklearn, вы можете создать mask для индексации NaN:

In [148]: mask = ~np.isnan(a)

In [149]: mean_squared_error(a[mask], b[mask])
Out[149]: 18.666666666666668