Merge pull request #16284 from eric-wieser/extract-lerp

MAINT: Clean up the implementation of quantile

1 files changed, 58 insertions, 62 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 0b23dbebd..cd3545966 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -3869,15 +3869,20 @@ def _quantile_is_valid(q):
     return True
 
 
+def _lerp(a, b, t, out=None):
+    """ Linearly interpolate from a to b by a factor of t """
+    return add(a*(1 - t), b*t, out=out)
+
+
 def _quantile_ureduce_func(a, q, axis=None, out=None, overwrite_input=False,
                            interpolation='linear', keepdims=False):
     a = asarray(a)
-    if q.ndim == 0:
-        # Do not allow 0-d arrays because following code fails for scalar
-        zerod = True
-        q = q[None]
-    else:
-        zerod = False
+
+    # ufuncs cause 0d array results to decay to scalars (see gh-13105), which
+    # makes them problematic for __setitem__ and attribute access. As a
+    # workaround, we call this on the result of every ufunc on a possibly-0d
+    # array.
+    not_scalar = np.asanyarray
 
     # prepare a for partitioning
     if overwrite_input:
@@ -3894,9 +3899,14 @@ def _quantile_ureduce_func(a, q, axis=None, out=None, overwrite_input=False,
     if axis is None:
         axis = 0
 
-    Nx = ap.shape[axis]
-    indices = q * (Nx - 1)
+    if q.ndim > 2:
+        # The code below works fine for nd, but it might not have useful
+        # semantics. For now, keep the supported dimensions the same as it was
+        # before.
+        raise ValueError("q must be a scalar or 1d")
 
+    Nx = ap.shape[axis]
+    indices = not_scalar(q * (Nx - 1))
     # round fractional indices according to interpolation method
     if interpolation == 'lower':
         indices = floor(indices).astype(intp)
@@ -3913,74 +3923,60 @@ def _quantile_ureduce_func(a, q, axis=None, out=None, overwrite_input=False,
             "interpolation can only be 'linear', 'lower' 'higher', "
             "'midpoint', or 'nearest'")
 
-    n = np.array(False, dtype=bool)  # check for nan's flag
-    if np.issubdtype(indices.dtype, np.integer):  # take the points along axis
-        # Check if the array contains any nan's
-        if np.issubdtype(a.dtype, np.inexact):
-            indices = concatenate((indices, [-1]))
-
-        ap.partition(indices, axis=axis)
-        # ensure axis with q-th is first
-        ap = np.moveaxis(ap, axis, 0)
-        axis = 0
-
-        # Check if the array contains any nan's
-        if np.issubdtype(a.dtype, np.inexact):
-            indices = indices[:-1]
-            n = np.isnan(ap[-1:, ...])
-
-        if zerod:
-            indices = indices[0]
-        r = take(ap, indices, axis=axis, out=out)
+    # The dimensions of `q` are prepended to the output shape, so we need the
+    # axis being sampled from `ap` to be first.
+    ap = np.moveaxis(ap, axis, 0)
+    del axis
 
-    else:  # weight the points above and below the indices
-        indices_below = floor(indices).astype(intp)
-        indices_above = indices_below + 1
-        indices_above[indices_above > Nx - 1] = Nx - 1
+    if np.issubdtype(indices.dtype, np.integer):
+        # take the points along axis
 
-        # Check if the array contains any nan's
         if np.issubdtype(a.dtype, np.inexact):
-            indices_above = concatenate((indices_above, [-1]))
+            # may contain nan, which would sort to the end
+            ap.partition(concatenate((indices.ravel(), [-1])), axis=0)
+            n = np.isnan(ap[-1])
+        else:
+            # cannot contain nan
+            ap.partition(indices.ravel(), axis=0)
+            n = np.array(False, dtype=bool)
 
-        ap.partition(concatenate((indices_below, indices_above)), axis=axis)
+        r = take(ap, indices, axis=0, out=out)
 
-        # ensure axis with q-th is first
-        ap = np.moveaxis(ap, axis, 0)
-        axis = 0
+    else:
+        # weight the points above and below the indices
 
-        weights_shape = [1] * ap.ndim
-        weights_shape[axis] = len(indices)
-        weights_above = (indices - indices_below).reshape(weights_shape)
+        indices_below = not_scalar(floor(indices)).astype(intp)
+        indices_above = not_scalar(indices_below + 1)
+        indices_above[indices_above > Nx - 1] = Nx - 1
 
-        # Check if the array contains any nan's
         if np.issubdtype(a.dtype, np.inexact):
-            indices_above = indices_above[:-1]
-            n = np.isnan(ap[-1:, ...])
+            # may contain nan, which would sort to the end
+            ap.partition(concatenate((
+                indices_below.ravel(), indices_above.ravel(), [-1]
+            )), axis=0)
+            n = np.isnan(ap[-1])
+        else:
+            # cannot contain nan
+            ap.partition(concatenate((
+                indices_below.ravel(), indices_above.ravel()
+            )), axis=0)
+            n = np.array(False, dtype=bool)
 
-        x1 = take(ap, indices_below, axis=axis) * (1 - weights_above)
-        x2 = take(ap, indices_above, axis=axis) * weights_above
+        weights_shape = indices.shape + (1,) * (ap.ndim - 1)
+        weights_above = not_scalar(indices - indices_below).reshape(weights_shape)
 
-        if zerod:
-            x1 = x1.squeeze(0)
-            x2 = x2.squeeze(0)
+        x_below = take(ap, indices_below, axis=0)
+        x_above = take(ap, indices_above, axis=0)
 
-        r = add(x1, x2, out=out)
+        r = _lerp(x_below, x_above, weights_above, out=out)
 
+    # if any slice contained a nan, then all results on that slice are also nan
     if np.any(n):
-        if zerod:
-            if ap.ndim == 1:
-                if out is not None:
-                    out[...] = a.dtype.type(np.nan)
-                    r = out
-                else:
-                    r = a.dtype.type(np.nan)
-            else:
-                r[..., n.squeeze(0)] = a.dtype.type(np.nan)
+        if r.ndim == 0 and out is None:
+            # can't write to a scalar
+            r = a.dtype.type(np.nan)
         else:
-            if r.ndim == 1:
-                r[:] = a.dtype.type(np.nan)
-            else:
-                r[..., n.repeat(q.size, 0)] = a.dtype.type(np.nan)
+            r[..., n] = a.dtype.type(np.nan)
 
     return r