From 4152443b1d268951401e0be1351eccbddd572e97 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Thu, 4 Feb 2021 18:32:54 +0100 Subject: MAINT: Avoid moveaxis overhead in median. This change speeds up taking the median of 1001 floats by ~20%, as measured by `python -mtimeit -s 'import numpy as np; x = np.random.randn(1001)' -- 'np.median(x)'` --- numpy/lib/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) (limited to 'numpy/lib/utils.py') diff --git a/numpy/lib/utils.py b/numpy/lib/utils.py index f7e176cf3..8f5c6eea3 100644 --- a/numpy/lib/utils.py +++ b/numpy/lib/utils.py @@ -1026,8 +1026,7 @@ def _median_nancheck(data, result, axis, out): """ if data.size == 0: return result - data = np.moveaxis(data, axis, -1) - n = np.isnan(data[..., -1]) + n = np.isnan(data.take(-1, axis=axis)) # masked NaN values are ok if np.ma.isMaskedArray(n): n = n.filled(False) -- cgit v1.2.1 From 3cb10bc7178b3a753d7cc1f81414a8ab94b92e64 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Tue, 9 Feb 2021 22:26:51 +0100 Subject: DOC: Fix docstring of _median_nancheck. _median_nancheck doesn't support axis being anything other than an integer (otherwise the call to moveaxis would fail). This is fine, because median goes through _ureduce for multi-axis support. --- numpy/lib/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'numpy/lib/utils.py') diff --git a/numpy/lib/utils.py b/numpy/lib/utils.py index f7e176cf3..9cf9c1e76 100644 --- a/numpy/lib/utils.py +++ b/numpy/lib/utils.py @@ -1015,8 +1015,8 @@ def _median_nancheck(data, result, axis, out): Input data to median function result : Array or MaskedArray Result of median function - axis : {int, sequence of int, None}, optional - Axis or axes along which the median was computed. + axis : int + Axis along which the median was computed. out : ndarray, optional Output array in which to place the result. Returns -- cgit v1.2.1