Repository: numpy/numpy

Issue:
ENH: rewrite ma.median to improve poor performance for multiple dimensions
many masked median along a small dimension is extremely slow due to the
usage of apply_along_axis which iterates fully in python. The unmasked
median is about 1000x faster.

Work around this issue by using indexing to select the median element
instead of apply_along_axis.

Further improvements are possible, e.g. using the current np.nanmedian
approach for masked medians along large dimensions so partition is used
instead of sort or to extend partition to allow broadcasting over
multiple elements.

Closes gh-4683.

np.ma.median returns masked_array; doesn't match example in docs
From the examples

```
    Examples
    --------
    >>> x = np.ma.array(np.arange(8), mask=[0]*4 + [1]*4)
    >>> np.ma.extras.median(x)
    1.5

```

On CentOS 6  64 with a conda python and numpy 1.9.2 

```
# python
Python 2.7.9 |Continuum Analytics, Inc.| (default, Mar  9 2015, 16:20:48) 
[GCC 4.4.7 20120313 (Red Hat 4.4.7-1)] on linux2
Type "help", "copyright", "credits" or "license" for more information.
Anaconda is brought to you by Continuum Analytics.
Please check out: http://continuum.io/thanks and https://binstar.org
>>> import numpy as np
>>> np.__version__
'1.9.2'
>>> x = np.ma.array(np.arange(8), mask=[0]*4 + [1]*4)
>>> np.ma.extras.median(x)
masked_array(data = [1.5],
             mask = [False],
       fill_value = 1e+20)
```

In addition to not matching the Example, I think this is also inconsistent with the return type text and with the behavior of other masked aggregate functions (np.ma.mean) which seems to just return a np.float64 on the same `x`.



Discussion:
I'd like this in 1.9 but it still needs some iterations, e.g. using partition instead of sort for larger dimensions.
this technique should also be used to the new nanmedian and nanpercentile.
posted now for comments, maybe I'm overlooking some neat indexing trick that would make this nicer.

added usage of masked median to nanmedian at a cutoff of 400 elements in the axis. on larger arrays the faster partition gains more than the apply_along_axis costs

as sorting nans has the same properties a masked sort (moving them to the end) one could probably duplicate the masked median code to squeeze some more performance as the masked sort is heavy on indexing and memory.
The masked median could probably also benefit from the threshold to use partition, but that involves quite extensive refactoring which would probably take too long to still go into 1.9.

I'm probably also not going to do the same for percentile for 1.9 anymore either.

Not sure of all the details, I'll leave that to the tests. 

How much gain does this bring to the small array performance? It might make sense to push this off to the next release so you can spend more time with it.

easily a factor 500, the performance of many small medians is abysmal as the pure python logic in apply_along_axis is very slow compared to the median.
I have seen many stumble over this issue, its pretty common in astronomy where you often stack a bunch of images using a median, e.g. 1000x1000x100 stacks. It also came up on the mailing list at least once.
the only existing workaround is using bottleneck or writing this yourself which not trivial.

Well, that's significant ;) Anything we could do to speed up apply_along_axis? 

in the linked issue I listed the ideas I had to solve this issue.
improving apply_along_axis should be possible and worthwhile by moving it to C code using npyiter, possibly one could also cythonize it

updated fixing the minor style issue and added an explicit testcase for the two nanmedian code paths

Thanks Julian.

That indeed looks inconsistent, would be good to make the behavior match the example again.

I think this was introduced by #4760 @juliantaylor.


Resolving patch:
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
--- 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
@@ -680,6 +680,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
@@ -688,10 +692,10 @@ def median(a, axis=None, out=None, overwrite_input=False):
     ind = np.meshgrid(*axes_grid, sparse=True, indexing='ij')
     # insert indices of low and high median
     ind.insert(axis, h - 1)
-    low = asorted[ind]
+    low = asorted[tuple(ind)]
     low._sharedmask = False
     ind[axis] = h
-    high = asorted[ind]
+    high = asorted[tuple(ind)]
     # duplicate high if odd number of elements so mean does nothing
     odd = counts % 2 == 1
     if asorted.ndim == 1:


Test patch:
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
--- 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)