BUG: nanpercentile/nanmedian 0-d with output given.

Also some PEP-8 fixes and test improvements

2 files changed, 95 insertions, 65 deletions

diff --git a/numpy/lib/nanfunctions.py b/numpy/lib/nanfunctions.py
index 739dcc3ac..478e7cf7e 100644
--- a/numpy/lib/nanfunctions.py
+++ b/numpy/lib/nanfunctions.py
@@ -603,11 +603,10 @@ def nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
     return avg
 
 
-def _nanmedian1d(arr1d, overwrite_input=False): # This only works on 1d arrays
+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]
@@ -617,10 +616,10 @@ def _nanmedian1d(arr1d, overwrite_input=False): # This only works on 1d arrays
     elif s.size == 0:
         return np.median(arr1d, overwrite_input=overwrite_input)
     else:
-        if overwrite_input: 
+        if overwrite_input:
             x = arr1d
-        else: 
-            x = arr1d.copy() 
+        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
@@ -641,12 +640,12 @@ def _nanmedian(a, axis=None, out=None, overwrite_input=False):
         if out is None:
             return _nanmedian1d(part, overwrite_input)
         else:
-            out[:] = _nanmedian1d(part, overwrite_input)
+            out[...] = _nanmedian1d(part, overwrite_input)
             return out
     else:
         result = np.apply_along_axis(_nanmedian1d, axis, a, overwrite_input)
         if out is not None:
-            out[:] = result
+            out[...] = result
         return result
 
 
@@ -688,9 +687,9 @@ def nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     Returns
     -------
     median : ndarray
-        A new array holding the result. If the input contains integers, or 
-        floats of smaller precision than 64, then the output data-type is 
-        float64.  Otherwise, the output data-type is the same as that of the 
+        A new array holding the result. If the input contains integers, or
+        floats of smaller precision than 64, then the output data-type is
+        float64.  Otherwise, the output data-type is the same as that of the
         input.
 
     See Also
@@ -732,7 +731,7 @@ def nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     a = np.asanyarray(a)
     # apply_along_axis in _nanmedian doesn't handle empty arrays well,
     # so deal them upfront
-    if 0 in a.shape:
+    if a.size == 0:
         return np.nanmean(a, axis, out=out, keepdims=keepdims)
 
     r, k = _ureduce(a, func=_nanmedian, axis=axis, out=out,
@@ -743,8 +742,8 @@ def nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False):
         return r
 
 
-def nanpercentile(a, q, axis=None, out=None,
-               overwrite_input=False, interpolation='linear', keepdims=False):
+def nanpercentile(a, q, axis=None, out=None, overwrite_input=False,
+                  interpolation='linear', keepdims=False):
     """
     Compute the qth percentile of the data along the specified axis, while
     ignoring nan values.
@@ -852,7 +851,7 @@ def nanpercentile(a, q, axis=None, out=None,
     q = np.asanyarray(q)
     # apply_along_axis in _nanpercentile doesn't handle empty arrays well,
     # so deal them upfront
-    if 0 in a.shape:
+    if a.size == 0:
         return np.nanmean(a, axis, out=out, keepdims=keepdims)
 
     r, k = _ureduce(a, func=_nanpercentile, q=q, axis=axis, out=out,
@@ -867,8 +866,8 @@ def nanpercentile(a, q, axis=None, out=None,
         return r
 
 
-def _nanpercentile(a, q, axis=None, out=None,
-                overwrite_input=False, interpolation='linear', keepdims=False):
+def _nanpercentile(a, q, axis=None, out=None, overwrite_input=False,
+                   interpolation='linear', keepdims=False):
     """
     Private function that doesn't support extended axis or keepdims.
     These methods are extended to this function using _ureduce
@@ -879,11 +878,11 @@ def _nanpercentile(a, q, axis=None, out=None,
         part = a.ravel()
         result = _nanpercentile1d(part, q, overwrite_input, interpolation)
     else:
-        result = np.apply_along_axis(_nanpercentile1d, axis, a, q, overwrite_input,
-                interpolation)
+        result = np.apply_along_axis(_nanpercentile1d, axis, a, q,
+                                     overwrite_input, interpolation)
 
     if out is not None:
-        out[:] = result
+        out[...] = result
     return result
 
 
@@ -891,7 +890,7 @@ def _nanpercentile1d(arr1d, q, overwrite_input=False, interpolation='linear'):
     """
     Private function for rank 1 arrays. Compute percentile ignoring NaNs.
     See nanpercentile for parameter usage
-    
+
     """
     c = np.isnan(arr1d)
     s = np.where(c)[0]
@@ -900,19 +899,19 @@ def _nanpercentile1d(arr1d, q, overwrite_input=False, interpolation='linear'):
         return np.nan
     elif s.size == 0:
         return np.percentile(arr1d, q, overwrite_input=overwrite_input,
-                interpolation=interpolation)
+                             interpolation=interpolation)
     else:
-        if overwrite_input: 
+        if overwrite_input:
             x = arr1d
-        else: 
-            x = arr1d.copy() 
+        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)
+                             interpolation=interpolation)
 
 
 def nanvar(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
diff --git a/numpy/lib/tests/test_nanfunctions.py b/numpy/lib/tests/test_nanfunctions.py
index 7f5217221..3fcfca218 100644
--- a/numpy/lib/tests/test_nanfunctions.py
+++ b/numpy/lib/tests/test_nanfunctions.py
@@ -542,11 +542,41 @@ class TestNanFunctions_Median(TestCase):
             res = np.nanmedian(mat, axis=axis, out=None, overwrite_input=False)
             assert_(res.ndim == tgt.ndim)
 
+        d = np.ones((3, 5, 7, 11))
+        # Randomly set some elements to NaN:
+        w = np.random.random((4, 200)) * np.array(d.shape)[:, None]
+        w = w.astype(np.intp)
+        d[tuple(w)] = np.nan
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always', RuntimeWarning)
+            res = np.nanmedian(d, axis=None, keepdims=True)
+            assert_equal(res.shape, (1, 1, 1, 1))
+            res = np.nanmedian(d, axis=(0, 1), keepdims=True)
+            assert_equal(res.shape, (1, 1, 7, 11))
+            res = np.nanmedian(d, axis=(0, 3), keepdims=True)
+            assert_equal(res.shape, (1, 5, 7, 1))
+            res = np.nanmedian(d, axis=(1,), keepdims=True)
+            assert_equal(res.shape, (3, 1, 7, 11))
+            res = np.nanmedian(d, axis=(0, 1, 2, 3), keepdims=True)
+            assert_equal(res.shape, (1, 1, 1, 1))
+            res = np.nanmedian(d, axis=(0, 1, 3), keepdims=True)
+            assert_equal(res.shape, (1, 1, 7, 1))
+
     def test_out(self):
-        mat = np.random.rand(3,3)
+        mat = np.random.rand(3, 3)
+        nan_mat = np.insert(mat, [0, 2], np.nan, axis=1)
         resout = np.zeros(3)
         tgt = np.median(mat, axis=1)
-        res = np.nanmedian(mat, axis=1, out=resout)
+        res = np.nanmedian(nan_mat, axis=1, out=resout)
+        assert_almost_equal(res, resout)
+        assert_almost_equal(res, tgt)
+        # 0-d output:
+        resout = np.zeros(())
+        tgt = np.median(mat, axis=None)
+        res = np.nanmedian(nan_mat, axis=None, out=resout)
+        assert_almost_equal(res, resout)
+        assert_almost_equal(res, tgt)
+        res = np.nanmedian(nan_mat, axis=(0, 1), out=resout)
         assert_almost_equal(res, resout)
         assert_almost_equal(res, tgt)
 
@@ -599,21 +629,6 @@ class TestNanFunctions_Median(TestCase):
         assert_raises(IndexError, np.nanmedian, d, axis=(0, 4))
         assert_raises(ValueError, np.nanmedian, d, axis=(1, 1))
 
-    def test_keepdims(self):
-        d = np.ones((3, 5, 7, 11))
-        assert_equal(np.nanmedian(d, axis=None, keepdims=True).shape,
-                     (1, 1, 1, 1))
-        assert_equal(np.nanmedian(d, axis=(0, 1), keepdims=True).shape,
-                     (1, 1, 7, 11))
-        assert_equal(np.nanmedian(d, axis=(0, 3), keepdims=True).shape,
-                     (1, 5, 7, 1))
-        assert_equal(np.nanmedian(d, axis=(1,), keepdims=True).shape,
-                     (3, 1, 7, 11))
-        assert_equal(np.nanmedian(d, axis=(0, 1, 2, 3), keepdims=True).shape,
-                     (1, 1, 1, 1))
-        assert_equal(np.nanmedian(d, axis=(0, 1, 3), keepdims=True).shape,
-                     (1, 1, 7, 1))
-
 
 class TestNanFunctions_Percentile(TestCase):
 
@@ -626,15 +641,47 @@ class TestNanFunctions_Percentile(TestCase):
     def test_keepdims(self):
         mat = np.eye(3)
         for axis in [None, 0, 1]:
-            tgt = np.percentile(mat, 70, axis=axis, out=None, overwrite_input=False)
-            res = np.percentile(mat, 70, axis=axis, out=None, overwrite_input=False)
+            tgt = np.percentile(mat, 70, axis=axis, out=None,
+                                overwrite_input=False)
+            res = np.nanpercentile(mat, 70, axis=axis, out=None,
+                                   overwrite_input=False)
             assert_(res.ndim == tgt.ndim)
 
+        d = np.ones((3, 5, 7, 11))
+        # Randomly set some elements to NaN:
+        w = np.random.random((4, 200)) * np.array(d.shape)[:, None]
+        w = w.astype(np.intp)
+        d[tuple(w)] = np.nan
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always', RuntimeWarning)
+            res = np.nanpercentile(d, 90, axis=None, keepdims=True)
+            assert_equal(res.shape, (1, 1, 1, 1))
+            res = np.nanpercentile(d, 90, axis=(0, 1), keepdims=True)
+            assert_equal(res.shape, (1, 1, 7, 11))
+            res = np.nanpercentile(d, 90, axis=(0, 3), keepdims=True)
+            assert_equal(res.shape, (1, 5, 7, 1))
+            res = np.nanpercentile(d, 90, axis=(1,), keepdims=True)
+            assert_equal(res.shape, (3, 1, 7, 11))
+            res = np.nanpercentile(d, 90, axis=(0, 1, 2, 3), keepdims=True)
+            assert_equal(res.shape, (1, 1, 1, 1))
+            res = np.nanpercentile(d, 90, axis=(0, 1, 3), keepdims=True)
+            assert_equal(res.shape, (1, 1, 7, 1))
+
     def test_out(self):
-        mat = np.random.rand(3,3)
+        mat = np.random.rand(3, 3)
+        nan_mat = np.insert(mat, [0, 2], np.nan, axis=1)
         resout = np.zeros(3)
         tgt = np.percentile(mat, 42, axis=1)
-        res = np.nanpercentile(mat, 42, axis=1, out=resout)
+        res = np.nanpercentile(nan_mat, 42, axis=1, out=resout)
+        assert_almost_equal(res, resout)
+        assert_almost_equal(res, tgt)
+        # 0-d output:
+        resout = np.zeros(())
+        tgt = np.percentile(mat, 42, axis=None)
+        res = np.nanpercentile(nan_mat, 42, axis=None, out=resout)
+        assert_almost_equal(res, resout)
+        assert_almost_equal(res, tgt)
+        res = np.nanpercentile(nan_mat, 42, axis=(0, 1), out=resout)
         assert_almost_equal(res, resout)
         assert_almost_equal(res, tgt)
 
@@ -642,10 +689,10 @@ class TestNanFunctions_Percentile(TestCase):
         tgt = [np.percentile(d, 28) for d in _rdat]
         res = np.nanpercentile(_ndat, 28, axis=1)
         assert_almost_equal(res, tgt)
-        tgt = [np.percentile(d, (28,98)) for d in _rdat]
-        res = np.nanpercentile(_ndat, (28,98), axis=1)
+        tgt = [np.percentile(d, (28, 98)) for d in _rdat]
+        res = np.nanpercentile(_ndat, (28, 98), axis=1)
         assert_almost_equal(res, tgt)
-    
+
     def test_allnans(self):
         mat = np.array([np.nan]*9).reshape(3, 3)
         for axis in [None, 0, 1]:
@@ -690,22 +737,6 @@ class TestNanFunctions_Percentile(TestCase):
         assert_raises(IndexError, np.nanpercentile, d, q=5, axis=(0, 4))
         assert_raises(ValueError, np.nanpercentile, d, q=5, axis=(1, 1))
 
-    def test_keepdims(self):
-        d = np.ones((3, 5, 7, 11))
-        assert_equal(np.nanpercentile(d, 90, axis=None, keepdims=True).shape,
-                     (1, 1, 1, 1))
-        assert_equal(np.nanpercentile(d, 90, axis=(0, 1), keepdims=True).shape,
-                     (1, 1, 7, 11))
-        assert_equal(np.nanpercentile(d, 90, axis=(0, 3), keepdims=True).shape,
-                     (1, 5, 7, 1))
-        assert_equal(np.nanpercentile(d, 90, axis=(1,), keepdims=True).shape,
-                     (3, 1, 7, 11))
-        assert_equal(np.nanpercentile(d, 90, axis=(0, 1, 2, 3), keepdims=True).shape,
-                     (1, 1, 1, 1))
-        assert_equal(np.nanpercentile(d, 90, axis=(0, 1, 3), keepdims=True).shape,
-                     (1, 1, 7, 1))
-
-
 
 if __name__ == "__main__":
     run_module_suite()