diff options
| author | Julian Taylor <jtaylor.debian@googlemail.com> | 2014-07-16 22:38:59 +0200 |
|---|---|---|
| committer | Julian Taylor <jtaylor.debian@googlemail.com> | 2014-07-16 22:49:34 +0200 |
| commit | acc3369adebf1cd5c6df7f58fe691be1a80a484a (patch) | |
| tree | 7e5ff72b090d69bfec4bc62cb7da88c647b60994 /numpy/ma/core.py | |
| parent | 10098daf387b9468a0aee19c3eb3e0cdd21f874c (diff) | |
| download | numpy-acc3369adebf1cd5c6df7f58fe691be1a80a484a.tar.gz | |
ENH: avoid meshgrid and fancy indexing for 1d masked sort
Improves performance by using the simple indexing path and memory by
avoiding creating a full meshgrid.
Diffstat (limited to 'numpy/ma/core.py')
| -rw-r--r-- | numpy/ma/core.py | 28 |
1 files changed, 20 insertions, 8 deletions
diff --git a/numpy/ma/core.py b/numpy/ma/core.py index 617f1921e..5c566b92c 100644 --- a/numpy/ma/core.py +++ b/numpy/ma/core.py @@ -5082,10 +5082,16 @@ class MaskedArray(ndarray): filler = maximum_fill_value(self) else: filler = fill_value - idx = np.meshgrid(*[np.arange(x) for x in self.shape], sparse=True, - indexing='ij') - idx[axis] = self.filled(filler).argsort(axis=axis, kind=kind, - order=order) + + sidx = self.filled(filler).argsort(axis=axis, kind=kind, + order=order) + # save meshgrid memory for 1d arrays + if self.ndim == 1: + idx = sidx + else: + idx = np.meshgrid(*[np.arange(x) for x in self.shape], sparse=True, + indexing='ij') + idx[axis] = sidx tmp_mask = self._mask[idx].flat tmp_data = self._data[idx].flat self._data.flat = tmp_data @@ -6187,10 +6193,16 @@ def sort(a, axis= -1, kind='quicksort', order=None, endwith=True, fill_value=Non filler = maximum_fill_value(a) else: filler = fill_value -# return - indx = np.meshgrid(*[np.arange(x) for x in a.shape], sparse=True, - indexing='ij') - indx[axis] = filled(a, filler).argsort(axis=axis, kind=kind, order=order) + + sindx = filled(a, filler).argsort(axis=axis, kind=kind, order=order) + + # save meshgrid memory for 1d arrays + if a.ndim == 1: + indx = sindx + else: + indx = np.meshgrid(*[np.arange(x) for x in a.shape], sparse=True, + indexing='ij') + indx[axis] = sindx return a[indx] sort.__doc__ = MaskedArray.sort.__doc__ |