From 6e5aa5e6b7980021ba5de67c872fad70dd4f9c92 Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Thu, 28 Apr 2016 19:05:27 -0400 Subject: [PATCH 1/2] TST: Verify np.ma.take works on scalars See #7585 --- numpy/ma/tests/test_core.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py index 217f307c6458..fc9ced365315 100644 --- a/numpy/ma/tests/test_core.py +++ b/numpy/ma/tests/test_core.py @@ -2998,6 +2998,10 @@ def test_take(self): assert_equal(x.take([[0, 1], [0, 1]]), masked_array([[10, 20], [10, 20]], [[0, 1], [0, 1]])) + # assert_equal crashes when passed np.ma.mask + self.assertIs(x[1], np.ma.masked) + self.assertIs(x.take(1), np.ma.masked) + x = array([[10, 20, 30], [40, 50, 60]], mask=[[0, 0, 1], [1, 0, 0, ]]) assert_equal(x.take([0, 2], axis=1), array([[10, 30], [40, 60]], mask=[[0, 1], [1, 0]])) From 090bbdd8f3debe0b6a6c4aa542937edf250f9443 Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Thu, 28 Apr 2016 19:07:33 -0400 Subject: [PATCH 2/2] BUG: Ensure mask is preserved on scalars By promoting and demoting between scalars and arrays where appropriate See #7585 --- numpy/ma/core.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/numpy/ma/core.py b/numpy/ma/core.py index e908a952cb9a..35c1ec791429 100644 --- a/numpy/ma/core.py +++ b/numpy/ma/core.py @@ -5610,9 +5610,10 @@ def take(self, indices, axis=None, out=None, mode='raise'): maskindices = getattr(indices, '_mask', nomask) if maskindices is not nomask: indices = indices.filled(0) - # Get the data + # Get the data, promoting scalars to 0d arrays with [...] so that + # .view works correctly if out is None: - out = _data.take(indices, axis=axis, mode=mode).view(cls) + out = _data.take(indices, axis=axis, mode=mode)[...].view(cls) else: np.take(_data, indices, axis=axis, mode=mode, out=out) # Get the mask @@ -5623,7 +5624,8 @@ def take(self, indices, axis=None, out=None, mode='raise'): outmask = _mask.take(indices, axis=axis, mode=mode) outmask |= maskindices out.__setmask__(outmask) - return out + # demote 0d arrays back to scalars, for consistency with ndarray.take + return out[()] # Array methods copy = _arraymethod('copy')