Merge pull request #7515 from ahaldane/ma_fix_negative_count

BUG: MaskedArray.count treats negative axes incorrectly

2 files changed, 12 insertions, 1 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 8593f62b2..e908a952c 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -4319,6 +4319,11 @@ class MaskedArray(ndarray):
                 return self.size
 
             axes = axis if isinstance(axis, tuple) else (axis,)
+            axes = tuple(a if a >= 0 else self.ndim + a for a in axes)
+            if len(axes) != len(set(axes)):
+                raise ValueError("duplicate value in 'axis'")
+            if np.any([a < 0 or a >= self.ndim for a in axes]):
+                raise ValueError("'axis' entry is out of bounds")
             items = 1
             for ax in axes:
                 items *= self.shape[ax]
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 6652e3969..217f307c6 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -991,7 +991,7 @@ class TestMaskedArrayArithmetic(TestCase):
         res = count(ott, 0)
         assert_(isinstance(res, ndarray))
         assert_(res.dtype.type is np.intp)
-        assert_raises(IndexError, ott.count, axis=1)
+        assert_raises(ValueError, ott.count, axis=1)
 
     def test_minmax_func(self):
         # Tests minimum and maximum.
@@ -4312,6 +4312,9 @@ class TestOptionalArgs(TestCase):
         assert_equal(count(a, keepdims=True), 16*ones((1,1,1)))
         assert_equal(count(a, axis=1, keepdims=True), 2*ones((2,1,4)))
         assert_equal(count(a, axis=(0,1), keepdims=True), 4*ones((1,1,4)))
+        assert_equal(count(a, axis=-2), 2*ones((2,4)))
+        assert_raises(ValueError, count, a, axis=(1,1))
+        assert_raises(ValueError, count, a, axis=3)
 
         # check the 'nomask' path
         a = np.ma.array(d, mask=nomask)
@@ -4322,6 +4325,9 @@ class TestOptionalArgs(TestCase):
         assert_equal(count(a, keepdims=True), 24*ones((1,1,1)))
         assert_equal(count(a, axis=1, keepdims=True), 3*ones((2,1,4)))
         assert_equal(count(a, axis=(0,1), keepdims=True), 6*ones((1,1,4)))
+        assert_equal(count(a, axis=-2), 3*ones((2,4)))
+        assert_raises(ValueError, count, a, axis=(1,1))
+        assert_raises(ValueError, count, a, axis=3)
 
         # check the 'masked' singleton
         assert_equal(count(np.ma.masked), 0)