Merge pull request #3769 from juliantaylor/percentile-compat

BUG: ensure percentile has same output structure as in 1.8

2 files changed, 182 insertions, 24 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 4d765bbc1..ae3d3d84c 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -2812,8 +2812,8 @@ def median(a, axis=None, out=None, overwrite_input=False):
     return mean(part[indexer], axis=axis, out=out)
 
 
-def percentile(a, q, interpolation='linear', axis=None, out=None,
-               overwrite_input=False):
+def percentile(a, q, axis=None, out=None,
+               overwrite_input=False, interpolation='linear'):
     """
     Compute the qth percentile of the data along the specified axis.
 
@@ -2825,15 +2825,6 @@ def percentile(a, q, interpolation='linear', axis=None, out=None,
         Input array or object that can be converted to an array.
     q : float in range of [0,100] (or sequence of floats)
         Percentile to compute which must be between 0 and 100 inclusive.
-    interpolation : {'linear', 'lower', 'higher', 'midpoint', 'nearest'}
-        This optional parameter specifies the interpolation method to use,
-        when the desired quantile lies between two data points `i` and `j`:
-            * linear: `i + (j - i) * fraction`, where `fraction` is the
-              fractional part of the index surrounded by `i` and `j`.
-            * lower: `i`.
-            * higher: `j`.
-            * nearest: `i` or `j` whichever is nearest.
-            * midpoint: (`i` + `j`) / 2.
     axis : int, optional
         Axis along which the percentiles are computed. The default (None)
         is to compute the percentiles along a flattened version of the array.
@@ -2851,15 +2842,28 @@ def percentile(a, q, interpolation='linear', axis=None, out=None,
         Note that, if the `a` input is not already an array this parameter
         will have no effect, `a` will be converted to an array internally
         regardless of the value of this parameter.
+    interpolation : {'linear', 'lower', 'higher', 'midpoint', 'nearest'}
+        This optional parameter specifies the interpolation method to use,
+        when the desired quantile lies between two data points `i` and `j`:
+            * linear: `i + (j - i) * fraction`, where `fraction` is the
+              fractional part of the index surrounded by `i` and `j`.
+            * lower: `i`.
+            * higher: `j`.
+            * nearest: `i` or `j` whichever is nearest.
+            * midpoint: (`i` + `j`) / 2.
+
+        .. versionadded:: 1.9.0
 
     Returns
     -------
-    percentile : ndarray
-        A new array holding the result (unless `out` is specified, in
-        which case that array is returned instead). 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.
+    percentile : scalar or ndarray
+        If a single percentile `q` is given and axis=None a scalar is returned.
+        If multiple percentiles `q` are given an array holding the result is
+        returned. The results are listed in the first axis.
+        (If `out` is specified, in which case that array is returned instead).
+        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
     --------
@@ -2906,7 +2910,14 @@ def percentile(a, q, interpolation='linear', axis=None, out=None,
 
     """
     a = asarray(a)
-    q = atleast_1d(q)
+    q = asarray(q)
+    if q.ndim == 0:
+        # Do not allow 0-d arrays because following code fails for scalar
+        zerod = True
+        q = q[None]
+    else:
+        zerod = False
+
     q = q / 100.0
     if (q < 0).any() or (q > 1).any():
         raise ValueError(
@@ -2948,7 +2959,13 @@ def percentile(a, q, interpolation='linear', axis=None, out=None,
 
     if indices.dtype == intp:  # take the points along axis
         ap.partition(indices, axis=axis)
-        return take(ap, indices, axis=axis, out=out)
+        # ensure axis with qth is first
+        ap = np.rollaxis(ap, axis, 0)
+        axis = 0
+
+        if zerod:
+            indices = indices[0]
+        r = take(ap, indices, axis=axis, out=out)
     else:  # weight the points above and below the indices
         indices_below = floor(indices).astype(intp)
         indices_above = indices_below + 1
@@ -2966,10 +2983,20 @@ def percentile(a, q, interpolation='linear', axis=None, out=None,
         x1 = take(ap, indices_below, axis=axis) * weights_below
         x2 = take(ap, indices_above, axis=axis) * weights_above
 
+        # ensure axis with qth is first
+        x1 = np.rollaxis(x1, axis, 0)
+        x2 = np.rollaxis(x2, axis, 0)
+
+        if zerod:
+            x1 = x1.squeeze(0)
+            x2 = x2.squeeze(0)
+
         if out is not None:
-            return add(x1, x2, out=out)
+            r = add(x1, x2, out=out)
         else:
-            return add(x1, x2)
+            r = add(x1, x2)
+
+    return r
 
 
 def trapz(y, x=None, dx=1.0, axis=-1):
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 8dd862d74..36f078de1 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1448,13 +1448,20 @@ class TestScoreatpercentile(TestCase):
         assert_equal(np.percentile(x, 100), 3.5)
         assert_equal(np.percentile(x, 50), 1.75)
 
+    def test_api(self):
+        d = np.ones(5)
+        np.percentile(d, 5, None, None, False)
+        np.percentile(d, 5, None, None, False, 'linear')
+        o = np.ones((1,))
+        np.percentile(d, 5, None, o, False, 'linear')
+
     def test_2D(self):
         x = np.array([[1, 1, 1],
                      [1, 1, 1],
                      [4, 4, 3],
                      [1, 1, 1],
                      [1, 1, 1]])
-        assert_array_equal(np.percentile(x, 50, axis=0), [[1, 1, 1]])
+        assert_array_equal(np.percentile(x, 50, axis=0), [1, 1, 1])
 
     def test_linear(self):
 
@@ -1496,7 +1503,79 @@ class TestScoreatpercentile(TestCase):
         assert_equal(np.percentile(x, (25, 50, 100), axis=0), r0)
 
         r1 = [[0.75, 1.5, 3], [4.75, 5.5, 7], [8.75, 9.5, 11]]
-        assert_equal(np.percentile(x, (25, 50, 100), axis=1), r1)
+        assert_equal(np.percentile(x, (25, 50, 100), axis=1), np.array(r1).T)
+
+        # ensure qth axis is always first as with np.array(old_percentile(..))
+        x = np.arange(3 * 4 * 5 * 6).reshape(3, 4, 5, 6)
+        assert_equal(np.percentile(x, (25, 50)).shape, (2,))
+        assert_equal(np.percentile(x, (25, 50, 75)).shape, (3,))
+        assert_equal(np.percentile(x, (25, 50), axis=0).shape, (2, 4, 5, 6))
+        assert_equal(np.percentile(x, (25, 50), axis=1).shape, (2, 3, 5, 6))
+        assert_equal(np.percentile(x, (25, 50), axis=2).shape, (2, 3, 4, 6))
+        assert_equal(np.percentile(x, (25, 50), axis=3).shape, (2, 3, 4, 5))
+        assert_equal(np.percentile(x, (25, 50, 75), axis=1).shape, (3, 3, 5, 6))
+        assert_equal(np.percentile(x, (25, 50),
+                                   interpolation="higher").shape, (2,))
+        assert_equal(np.percentile(x, (25, 50, 75),
+                                   interpolation="higher").shape, (3,))
+        assert_equal(np.percentile(x, (25, 50), axis=0,
+                                   interpolation="higher").shape, (2, 4, 5, 6))
+        assert_equal(np.percentile(x, (25, 50), axis=1,
+                                   interpolation="higher").shape, (2, 3, 5, 6))
+        assert_equal(np.percentile(x, (25, 50), axis=2,
+                                   interpolation="higher").shape, (2, 3, 4, 6))
+        assert_equal(np.percentile(x, (25, 50), axis=3,
+                                   interpolation="higher").shape, (2, 3, 4, 5))
+        assert_equal(np.percentile(x, (25, 50, 75), axis=1,
+                                   interpolation="higher").shape, (3, 3, 5, 6))
+
+    def test_scalar_q(self):
+        # test for no empty dimensions for compatiblity with old percentile
+        x = np.arange(12).reshape(3, 4)
+        assert_equal(np.percentile(x, 50), 5.5)
+        self.assertTrue(np.isscalar(np.percentile(x, 50)))
+        r0 = np.array([ 4.,  5.,  6.,  7.])
+        assert_equal(np.percentile(x, 50, axis=0), r0)
+        assert_equal(np.percentile(x, 50, axis=0).shape, r0.shape)
+        r1 = np.array([ 1.5,  5.5,  9.5])
+        assert_almost_equal(np.percentile(x, 50, axis=1), r1)
+        assert_equal(np.percentile(x, 50, axis=1).shape, r1.shape)
+
+        out = np.empty(1)
+        assert_equal(np.percentile(x, 50, out=out), 5.5)
+        assert_equal(out, 5.5)
+        out = np.empty(4)
+        assert_equal(np.percentile(x, 50, axis=0, out=out), r0)
+        assert_equal(out, r0)
+        out = np.empty(3)
+        assert_equal(np.percentile(x, 50, axis=1, out=out), r1)
+        assert_equal(out, r1)
+
+        # test for no empty dimensions for compatiblity with old percentile
+        x = np.arange(12).reshape(3, 4)
+        assert_equal(np.percentile(x, 50, interpolation='lower'), 5.)
+        self.assertTrue(np.isscalar(np.percentile(x, 50)))
+        r0 = np.array([ 4.,  5.,  6.,  7.])
+        c0 = np.percentile(x, 50, interpolation='lower', axis=0)
+        assert_equal(c0, r0)
+        assert_equal(c0.shape, r0.shape)
+        r1 = np.array([ 1.,  5.,  9.])
+        c1 = np.percentile(x, 50, interpolation='lower', axis=1)
+        assert_almost_equal(c1, r1)
+        assert_equal(c1.shape, r1.shape)
+
+        out = np.empty((), dtype=x.dtype)
+        c = np.percentile(x, 50, interpolation='lower', out=out)
+        assert_equal(c, 5)
+        assert_equal(out, 5)
+        out = np.empty(4, dtype=x.dtype)
+        c = np.percentile(x, 50, interpolation='lower', axis=0, out=out)
+        assert_equal(c, r0)
+        assert_equal(out, r0)
+        out = np.empty(3, dtype=x.dtype)
+        c = np.percentile(x, 50, interpolation='lower', axis=1, out=out)
+        assert_equal(c, r1)
+        assert_equal(out, r1)
 
     def test_exception(self):
         assert_raises(ValueError, np.percentile, [1, 2], 56,
@@ -1521,10 +1600,62 @@ class TestScoreatpercentile(TestCase):
         np.percentile(x, p, axis=0, out=y)
         assert_equal(y, np.percentile(x, p, axis=0))
 
-        y = np.zeros((2, 3))
+        y = np.zeros((3, 2))
         np.percentile(x, p, axis=1, out=y)
         assert_equal(y, np.percentile(x, p, axis=1))
 
+        x = np.arange(12).reshape(3, 4)
+        # q.dim > 1, float
+        r0 = np.array([[2.,  3.,  4., 5.], [4., 5., 6., 7.]])
+        out = np.empty((2, 4))
+        assert_equal(np.percentile(x, (25, 50), axis=0, out=out), r0)
+        assert_equal(out, r0)
+        r1 = np.array([[0.75,  4.75,  8.75], [1.5,  5.5,  9.5]])
+        out = np.empty((2, 3))
+        assert_equal(np.percentile(x, (25, 50), axis=1, out=out), r1)
+        assert_equal(out, r1)
+
+        # q.dim > 1, int
+        r0 = np.array([[0,  1,  2, 3], [4, 5, 6, 7]])
+        out = np.empty((2, 4), dtype=x.dtype)
+        c = np.percentile(x, (25, 50), interpolation='lower', axis=0, out=out)
+        assert_equal(c, r0)
+        assert_equal(out, r0)
+        r1 = np.array([[0,  4,  8], [1,  5,  9]])
+        out = np.empty((2, 3), dtype=x.dtype)
+        c = np.percentile(x, (25, 50), interpolation='lower', axis=1, out=out)
+        assert_equal(c, r1)
+        assert_equal(out, r1)
+
+    def test_percentile_empty_dim(self):
+        # empty dims are preserved
+        d = np.arange(11*2).reshape(11, 1, 2, 1)
+        assert_array_equal(np.percentile(d, 50, axis=0).shape, (1, 2, 1))
+        assert_array_equal(np.percentile(d, 50, axis=1).shape, (11, 2, 1))
+        assert_array_equal(np.percentile(d, 50, axis=2).shape, (11, 1, 1))
+        assert_array_equal(np.percentile(d, 50, axis=3).shape, (11, 1, 2))
+        assert_array_equal(np.percentile(d, 50, axis=-1).shape, (11, 1, 2))
+        assert_array_equal(np.percentile(d, 50, axis=-2).shape, (11, 1, 1))
+        assert_array_equal(np.percentile(d, 50, axis=-3).shape, (11, 2, 1))
+        assert_array_equal(np.percentile(d, 50, axis=-4).shape, (1, 2, 1))
+
+        assert_array_equal(np.percentile(d, 50, axis=2,
+                                         interpolation='midpoint').shape,
+                           (11, 1, 1))
+        assert_array_equal(np.percentile(d, 50, axis=-2,
+                                         interpolation='midpoint').shape,
+                           (11, 1, 1))
+
+        assert_array_equal(np.array(np.percentile(d, [10, 50], axis=0)).shape,
+                           (2, 1, 2, 1))
+        assert_array_equal(np.array(np.percentile(d, [10, 50], axis=1)).shape,
+                           (2, 11, 2, 1))
+        assert_array_equal(np.array(np.percentile(d, [10, 50], axis=2)).shape,
+                           (2, 11, 1, 1))
+        assert_array_equal(np.array(np.percentile(d, [10, 50], axis=3)).shape,
+                           (2, 11, 1, 2))
+
+
     def test_percentile_no_overwrite(self):
         a = np.array([2, 3, 4, 1])
         np.percentile(a, [50], overwrite_input=False)