diff options
| -rw-r--r-- | numpy/lib/nanfunctions.py | 27 | ||||
| -rw-r--r-- | numpy/lib/tests/test_nanfunctions.py | 28 |
2 files changed, 55 insertions, 0 deletions
diff --git a/numpy/lib/nanfunctions.py b/numpy/lib/nanfunctions.py index b3bf1880b..77c851fcf 100644 --- a/numpy/lib/nanfunctions.py +++ b/numpy/lib/nanfunctions.py @@ -40,6 +40,33 @@ __all__ = [ ] +def _nan_mask(a, out=None): + """ + Parameters + ---------- + a : array-like + Input array with at least 1 dimension. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output and will prevent the allocation of a new array. + + Returns + ------- + y : bool ndarray or True + A bool array where ``np.nan`` positions are marked with ``False`` + and other positions are marked with ``True``. If the type of ``a`` + is such that it can't possibly contain ``np.nan``, returns ``True``. + """ + # we assume that a is an array for this private function + + if a.dtype.kind not in 'fc': + return True + + y = np.isnan(a, out=out) + y = np.invert(y, out=y) + return y + def _replace_nan(a, val): """ If `a` is of inexact type, make a copy of `a`, replace NaNs with diff --git a/numpy/lib/tests/test_nanfunctions.py b/numpy/lib/tests/test_nanfunctions.py index 504372faf..b7261c63f 100644 --- a/numpy/lib/tests/test_nanfunctions.py +++ b/numpy/lib/tests/test_nanfunctions.py @@ -1,8 +1,10 @@ from __future__ import division, absolute_import, print_function import warnings +import pytest import numpy as np +from numpy.lib.nanfunctions import _nan_mask from numpy.testing import ( assert_, assert_equal, assert_almost_equal, assert_no_warnings, assert_raises, assert_array_equal, suppress_warnings @@ -925,3 +927,29 @@ class TestNanFunctions_Quantile(object): p = p.tolist() np.nanquantile(np.arange(100.), p, interpolation="midpoint") assert_array_equal(p, p0) + +@pytest.mark.parametrize("arr, expected", [ + # array of floats with some nans + (np.array([np.nan, 5.0, np.nan, np.inf]), + np.array([False, True, False, True])), + # int64 array that can't possibly have nans + (np.array([1, 5, 7, 9], dtype=np.int64), + True), + # bool array that can't possibly have nans + (np.array([False, True, False, True]), + True), + # 2-D complex array with nans + (np.array([[np.nan, 5.0], + [np.nan, np.inf]], dtype=np.complex64), + np.array([[False, True], + [False, True]])), + ]) +def test__nan_mask(arr, expected): + for out in [None, np.empty(arr.shape, dtype=np.bool_)]: + actual = _nan_mask(arr, out=out) + assert_equal(actual, expected) + # the above won't distinguish between True proper + # and an array of True values; we want True proper + # for types that can't possibly contain NaN + if type(expected) is not np.ndarray: + assert actual is True |