summaryrefslogtreecommitdiff
path: root/numpy/lib/function_base.py
diff options
context:
space:
mode:
authorEric Wieser <wieser.eric@gmail.com>2020-05-17 15:14:58 +0100
committerEric Wieser <wieser.eric@gmail.com>2020-05-17 15:14:58 +0100
commit49a84f6fd548c0d01d5dbd6ee286cdaa49fd6209 (patch)
treee212b76f99a587822c99a65c6ac8898e26db103c /numpy/lib/function_base.py
parentc2cbf13ac4df693cec70f6949efdc386c791d1f9 (diff)
downloadnumpy-49a84f6fd548c0d01d5dbd6ee286cdaa49fd6209.tar.gz
MAINT: Avoid moving axes around multiple times
It's easier to move the relevant axis to position 0 in `ap` first than it is to move it for every relevant object simultaneously.
Diffstat (limited to 'numpy/lib/function_base.py')
-rw-r--r--numpy/lib/function_base.py20
1 files changed, 5 insertions, 15 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 12ba3b1c6..2d1adc362 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -3943,34 +3943,24 @@ def _quantile_ureduce_func(a, q, axis=None, out=None, overwrite_input=False,
if np.issubdtype(a.dtype, np.inexact):
indices_above = concatenate((indices_above, [-1]))
- weights_above = indices - indices_below
- weights_below = 1 - weights_above
-
- weights_shape = [1, ] * ap.ndim
- weights_shape[axis] = len(indices)
- weights_below.shape = weights_shape
- weights_above.shape = weights_shape
-
ap.partition(concatenate((indices_below, indices_above)), axis=axis)
# ensure axis with q-th is first
ap = np.moveaxis(ap, axis, 0)
- weights_below = np.moveaxis(weights_below, axis, 0)
- weights_above = np.moveaxis(weights_above, axis, 0)
axis = 0
+ weights_shape = [1] * ap.ndim
+ weights_shape[axis] = len(indices)
+ weights_above = (indices - indices_below).reshape(weights_shape)
+
# 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:, ...])
- x1 = take(ap, indices_below, axis=axis) * weights_below
+ x1 = take(ap, indices_below, axis=axis) * (1 - weights_above)
x2 = take(ap, indices_above, axis=axis) * weights_above
- # ensure axis with q-th is first
- x1 = np.moveaxis(x1, axis, 0)
- x2 = np.moveaxis(x2, axis, 0)
-
if zerod:
x1 = x1.squeeze(0)
x2 = x2.squeeze(0)