def test__replace_nan(): """ Test that _replace_nan returns the original array if there are no NaNs, not a copy. """ for dtype in [np.bool, np.int32, np.int64]: arr = np.array([0, 1], dtype=dtype) result, mask = _replace_nan(arr, 0) assert mask is None # do not make a copy if there are no nans assert result is arr for dtype in [np.float32, np.float64]: arr = np.array([0, 1], dtype=dtype) result, mask = _replace_nan(arr, 2) assert (mask == False).all() # mask is not None, so we make a copy assert result is not arr assert_equal(result, arr) arr_nan = np.array([0, 1, np.nan], dtype=dtype) result_nan, mask_nan = _replace_nan(arr_nan, 2) assert_equal(mask_nan, np.array([False, False, True])) assert result_nan is not arr_nan assert_equal(result_nan, np.array([0, 1, 2])) assert np.isnan(arr_nan[-1])
def nanargmin(a, axis): from numpy.lib.nanfunctions import _replace_nan a, mask = _replace_nan(a, np.inf) res = np.argmin(a, axis=axis) return res