Repository: numba/numba

Issue:
Implement support for various numpy reduce functions
Implement support for various numpy reduce functions like sum, mean, median, etc.



Discussion:
This was mostly done in PR #833. median is still unsupported.

@pitrou Can I take up the work on implementing median?

Yes, you can!


Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -13,7 +13,7 @@
 import numba.ctypes_support as ctypes
 import numpy
 from llvmlite.llvmpy.core import Constant
-from numba import types, cgutils
+from numba import types, cgutils, typing
 from numba.numpy_support import as_dtype
 from numba.numpy_support import version as numpy_version
 from numba.targets.imputils import (builtin, builtin_attr, implement,
@@ -904,6 +904,70 @@ def array_argmax_impl(arry):
     return context.compile_internal(builder, array_argmax_impl, sig, args)
 
 
+@builtin
+@implement(numpy.median, types.Kind(types.Array))
+def array_median(context, builder, sig, args):
+
+    def partition(A, low, high):
+        mid = (low+high) // 2
+        # median of three {low, middle, high}
+        LM = A[low] <= A[mid]
+        MH = A[mid] <= A[high]
+        LH = A[low] <= A[high]
+
+        if LM == MH:
+            median3 = mid
+        elif LH != LM:
+            median3 = low
+        else:
+            median3 = high
+
+        # choose median3 as the pivot
+        A[high], A[median3] = A[median3], A[high]
+
+        x = A[high]
+        i = low
+        for j in range(low, high):
+            if A[j] <= x:
+                A[i], A[j] = A[j], A[i]
+                i += 1
+        A[i], A[high] = A[high], A[i]
+        return i
+
+    sig_partition = typing.signature(types.intp, *(sig.args[0], types.intp, types.intp))
+    _partition = context.compile_subroutine(builder, partition, sig_partition)
+
+    def select(arry, k):
+        n = arry.shape[0]
+        # XXX: assuming flat array till array.flatten is implemented
+        # temp_arry = arry.flatten()
+        temp_arry = arry.copy()
+        high = n-1
+        low = 0
+        # NOTE: high is inclusive
+        i = _partition(temp_arry, low, high)
+        while i != k:
+            if i < k:
+                low = i+1
+                i = _partition(temp_arry, low, high)
+            else:
+                high = i-1
+                i = _partition(temp_arry, low, high)
+        return temp_arry[k]
+
+    sig_select = typing.signature(sig.args[0].dtype, *(sig.args[0], types.intp))
+    _select = context.compile_subroutine(builder, select, sig_select)
+
+    def median(arry):
+        n = arry.shape[0]
+        if n % 2 == 0:
+            return (_select(arry, n//2 - 1) + _select(arry, n//2))/2
+        else:
+            return _select(arry, n//2)
+
+    return context.compile_internal(builder, median, sig, args)
+
+
 def _np_round_intrinsic(tp):
     # np.round() always rounds half to even
     return "llvm.rint.f%d" % (tp.bitwidth,)
diff --git a/numba/typing/builtins.py b/numba/typing/builtins.py
--- a/numba/typing/builtins.py
+++ b/numba/typing/builtins.py
@@ -714,7 +714,7 @@ def array_attribute_attachment(self, ary):
     install_array_method(fname, generic_expand_cumulative)
 
 # Functions that require integer arrays get promoted to float64 return
-for fName in ["mean", "var", "std"]:
+for fName in ["mean", "median", "var", "std"]:
     install_array_method(fName, generic_hetero_real)
 
 # Functions that return an index (intp)
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -332,7 +332,7 @@ def _numpy_redirect(fname):
                dict(key=numpy_function, method_name=fname))
     builtin_global(numpy_function, types.Function(cls))
 
-for func in ['min', 'max', 'sum', 'prod', 'mean', 'var', 'std',
+for func in ['min', 'max', 'sum', 'prod', 'mean', 'median', 'var', 'std',
              'cumsum', 'cumprod', 'argmin', 'argmax']:
     _numpy_redirect(func)
 


Test patch:
diff --git a/numba/tests/test_array_reductions.py b/numba/tests/test_array_reductions.py
--- a/numba/tests/test_array_reductions.py
+++ b/numba/tests/test_array_reductions.py
@@ -77,6 +77,9 @@ def array_argmax(arr):
 def array_argmax_global(arr):
     return np.argmax(arr)
 
+def array_median_global(arr):
+    return np.median(arr)
+
 
 def base_test_arrays(dtype):
     a1 = np.arange(10, dtype=dtype) + 1
@@ -148,6 +151,33 @@ def test_argmax_basic(self):
         npr, nbr = run_comparative(array_argmax, arr)
         self.assertPreciseEqual(npr, nbr)
 
+    def test_median_odd(self):
+        arr = np.arange(101)
+        np.random.shuffle(arr)
+        npr, nbr = run_comparative(array_median_global, arr)
+        self.assertPreciseEqual(npr, nbr)
+
+    def test_median_even(self):
+        arr = np.arange(100)
+        np.random.shuffle(arr)
+        npr, nbr = run_comparative(array_median_global, arr)
+        self.assertPreciseEqual(npr, nbr)
+
+    def test_median_sorted(self):
+        arr = np.arange(100)
+        npr, nbr = run_comparative(array_median_global, arr)
+        self.assertPreciseEqual(npr, nbr)
+
+    def test_median_revsorted(self):
+        arr = np.arange(100, 0, -1)
+        npr, nbr = run_comparative(array_median_global, arr)
+        self.assertPreciseEqual(npr, nbr)
+
+    def test_median_duplicate(self):
+        arr = np.ones(100)
+        npr, nbr = run_comparative(array_median_global, arr)
+        self.assertPreciseEqual(npr, nbr)
+
     def test_array_sum_global(self):
         arr = np.arange(10, dtype=np.int32)
         arrty = typeof(arr)