BUG: ma.median of 1d array should return a scalar

AmitAronovitch · AmitAronovitch · commit bb46a495ee52 · 2016-05-22T00:04:21.000+03:00
Fixes #5969. Performance fix #4760 had caused wrong shaped results in the 1D case. This fix restores the original 1D behavior.
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
@@ -27,7 +27,7 @@
 
 from . import core as ma
 from .core import (
-    MaskedArray, MAError, add, array, asarray, concatenate, filled,
+    MaskedArray, MAError, add, array, asarray, concatenate, filled, count,
     getmask, getmaskarray, make_mask_descr, masked, masked_array, mask_or,
     nomask, ones, sort, zeros, getdata, get_masked_subclass, dot,
     mask_rowcols
@@ -653,6 +653,10 @@ def _median(a, axis=None, out=None, overwrite_input=False):
     elif axis < 0:
         axis += a.ndim
 
+    if asorted.ndim == 1:
+        idx, odd = divmod(count(asorted), 2)
+        return asorted[idx - (not odd) : idx + 1].mean()
+
     counts = asorted.shape[axis] - (asorted.mask).sum(axis=axis)
     h = counts // 2
     # create indexing mesh grid for all but reduced axis
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
@@ -662,6 +662,19 @@ def test_non_masked(self):
         assert_equal(np.ma.median(np.arange(9)), 4.)
         assert_equal(np.ma.median(range(9)), 4)
 
+    def test_masked_1d(self):
+        "test the examples given in the docstring of ma.median"
+        x = array(np.arange(8), mask=[0]*4 + [1]*4)
+        assert_equal(np.ma.median(x), 1.5)
+        assert_equal(np.ma.median(x).shape, (), "shape mismatch")
+        x = array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4)
+        assert_equal(np.ma.median(x), 2.5)
+        assert_equal(np.ma.median(x).shape, (), "shape mismatch")
+
+    def test_1d_shape_consistency(self):
+        assert_equal(np.ma.median(array([1,2,3],mask=[0,0,0])).shape,
+                     np.ma.median(array([1,2,3],mask=[0,1,0])).shape )
+
     def test_2d(self):
         # Tests median w/ 2D
         (n, p) = (101, 30)
