diff options
Diffstat (limited to 'numpy/lib/nanfunctions.py')
| -rw-r--r-- | numpy/lib/nanfunctions.py | 105 |
1 files changed, 56 insertions, 49 deletions
diff --git a/numpy/lib/nanfunctions.py b/numpy/lib/nanfunctions.py index 1e342b932..79bf01281 100644 --- a/numpy/lib/nanfunctions.py +++ b/numpy/lib/nanfunctions.py @@ -106,6 +106,46 @@ def _copyto(a, val, mask): return a +def _remove_nan_1d(arr1d, overwrite_input=False): + """ + Equivalent to arr1d[~arr1d.isnan()], but in a different order + + Presumably faster as it incurs fewer copies + + Parameters + ---------- + arr1d : ndarray + Array to remove nans from + overwrite_input : bool + True if `arr1d` can be modified in place + + Returns + ------- + res : ndarray + Array with nan elements removed + overwrite_input : bool + True if `res` can be modified in place, given the constraint on the + input + """ + + c = np.isnan(arr1d) + s = np.nonzero(c)[0] + if s.size == arr1d.size: + warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=4) + return arr1d[:0], True + elif s.size == 0: + return arr1d, overwrite_input + else: + if not overwrite_input: + arr1d = arr1d.copy() + # select non-nans at end of array + enonan = arr1d[-s.size:][~c[-s.size:]] + # fill nans in beginning of array with non-nans of end + arr1d[s[:enonan.size]] = enonan + + return arr1d[:-s.size], True + + def _divide_by_count(a, b, out=None): """ Compute a/b ignoring invalid results. If `a` is an array the division @@ -836,24 +876,12 @@ def _nanmedian1d(arr1d, overwrite_input=False): Private function for rank 1 arrays. Compute the median ignoring NaNs. See nanmedian for parameter usage """ - c = np.isnan(arr1d) - s = np.where(c)[0] - if s.size == arr1d.size: - warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=3) + arr1d, overwrite_input = _remove_nan_1d(arr1d, + overwrite_input=overwrite_input) + if arr1d.size == 0: return np.nan - elif s.size == 0: - return np.median(arr1d, overwrite_input=overwrite_input) - else: - if overwrite_input: - x = arr1d - else: - x = arr1d.copy() - # select non-nans at end of array - enonan = arr1d[-s.size:][~c[-s.size:]] - # fill nans in beginning of array with non-nans of end - x[s[:enonan.size]] = enonan - # slice nans away - return np.median(x[:-s.size], overwrite_input=True) + + return np.median(arr1d, overwrite_input=overwrite_input) def _nanmedian(a, axis=None, out=None, overwrite_input=False): @@ -1088,7 +1116,7 @@ def nanpercentile(a, q, axis=None, out=None, overwrite_input=False, >>> a[0][1] = np.nan >>> a array([[ 10., nan, 4.], - [ 3., 2., 1.]]) + [ 3., 2., 1.]]) >>> np.percentile(a, 50) nan >>> np.nanpercentile(a, 50) @@ -1123,10 +1151,7 @@ def nanpercentile(a, q, axis=None, out=None, overwrite_input=False, overwrite_input=overwrite_input, interpolation=interpolation) if keepdims and keepdims is not np._NoValue: - if q.ndim == 0: - return r.reshape(k) - else: - return r.reshape([len(q)] + k) + return r.reshape(q.shape + k) else: return r @@ -1149,7 +1174,7 @@ def _nanpercentile(a, q, axis=None, out=None, overwrite_input=False, # Move that axis to the beginning to match percentile's # convention. if q.ndim != 0: - result = np.rollaxis(result, axis) + result = np.moveaxis(result, axis, 0) if out is not None: out[...] = result @@ -1158,34 +1183,16 @@ def _nanpercentile(a, q, axis=None, out=None, overwrite_input=False, def _nanpercentile1d(arr1d, q, overwrite_input=False, interpolation='linear'): """ - Private function for rank 1 arrays. Compute percentile ignoring - NaNs. - + Private function for rank 1 arrays. Compute percentile ignoring NaNs. See nanpercentile for parameter usage """ - c = np.isnan(arr1d) - s = np.where(c)[0] - if s.size == arr1d.size: - warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=3) - if q.ndim == 0: - return np.nan - else: - return np.nan * np.ones((len(q),)) - elif s.size == 0: - return np.percentile(arr1d, q, overwrite_input=overwrite_input, - interpolation=interpolation) - else: - if overwrite_input: - x = arr1d - else: - x = arr1d.copy() - # select non-nans at end of array - enonan = arr1d[-s.size:][~c[-s.size:]] - # fill nans in beginning of array with non-nans of end - x[s[:enonan.size]] = enonan - # slice nans away - return np.percentile(x[:-s.size], q, overwrite_input=True, - interpolation=interpolation) + arr1d, overwrite_input = _remove_nan_1d(arr1d, + overwrite_input=overwrite_input) + if arr1d.size == 0: + return np.full(q.shape, np.nan)[()] # convert to scalar + + return np.percentile(arr1d, q, overwrite_input=overwrite_input, + interpolation=interpolation) def nanvar(a, axis=None, dtype=None, out=None, ddof=0, keepdims=np._NoValue): |