Merge pull request #7823 from madphysicist/ma-bugs

BUG: Fixed masked array behavior for scalar inputs to np.ma.atleast_*d

4 files changed, 133 insertions, 52 deletions

diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
index 31d7d1c03..2b5110473 100644
--- a/numpy/core/shape_base.py
+++ b/numpy/core/shape_base.py
@@ -146,7 +146,7 @@ def atleast_3d(*arys):
     >>> x = np.arange(12.0).reshape(4,3)
     >>> np.atleast_3d(x).shape
     (4, 3, 1)
-    >>> np.atleast_3d(x).base is x
+    >>> np.atleast_3d(x).base is x.base  # x is a reshape, so not base itself
     True
 
     >>> for arr in np.atleast_3d([1, 2], [[1, 2]], [[[1, 2]]]):
diff --git a/numpy/lib/info.py b/numpy/lib/info.py
index ca1e72397..141df2ace 100644
--- a/numpy/lib/info.py
+++ b/numpy/lib/info.py
@@ -67,9 +67,9 @@ Shape Manipulation
 ------------------
 ================ ===================
 squeeze          Return a with length-one dimensions removed.
-atleast_1d       Force arrays to be > 1D
-atleast_2d       Force arrays to be > 2D
-atleast_3d       Force arrays to be > 3D
+atleast_1d       Force arrays to be >= 1D
+atleast_2d       Force arrays to be >= 2D
+atleast_3d       Force arrays to be >= 3D
 vstack           Stack arrays vertically (row on row)
 hstack           Stack arrays horizontally (column on column)
 column_stack     Stack 1D arrays as columns into 2D array
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index cbf7b6cdb..aefdd4e49 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -45,9 +45,8 @@ def issequence(seq):
     Is seq a sequence (ndarray, list or tuple)?
 
     """
-    if isinstance(seq, (ndarray, tuple, list)):
-        return True
-    return False
+    return isinstance(seq, (ndarray, tuple, list))
+
 
 def count_masked(arr, axis=None):
     """
@@ -103,6 +102,7 @@ def count_masked(arr, axis=None):
     m = getmaskarray(arr)
     return m.sum(axis)
 
+
 def masked_all(shape, dtype=float):
     """
     Empty masked array with all elements masked.
@@ -154,6 +154,7 @@ def masked_all(shape, dtype=float):
                      mask=np.ones(shape, make_mask_descr(dtype)))
     return a
 
+
 def masked_all_like(arr):
     """
     Empty masked array with the properties of an existing array.
@@ -221,6 +222,9 @@ class _fromnxfunction:
     as the wrapped NumPy function. The docstring of `newfunc` is adapted from
     the wrapped function as well, see `getdoc`.
 
+    This class should not be used directly. Instead, one of its extensions that
+    provides support for a specific type of input should be used.
+
     Parameters
     ----------
     funcname : str
@@ -261,48 +265,100 @@ class _fromnxfunction:
         return
 
     def __call__(self, *args, **params):
+        pass
+
+
+class _fromnxfunction_single(_fromnxfunction):
+    """
+    A version of `_fromnxfunction` that is called with a single array
+    argument followed by auxiliary args that are passed verbatim for
+    both the data and mask calls.
+    """
+    def __call__(self, x, *args, **params):
         func = getattr(np, self.__name__)
-        if len(args) == 1:
-            x = args[0]
-            if isinstance(x, ndarray):
-                _d = func(x.__array__(), **params)
-                _m = func(getmaskarray(x), **params)
-                return masked_array(_d, mask=_m)
-            elif isinstance(x, tuple) or isinstance(x, list):
-                _d = func(tuple([np.asarray(a) for a in x]), **params)
-                _m = func(tuple([getmaskarray(a) for a in x]), **params)
-                return masked_array(_d, mask=_m)
-            else:
-                _d = func(np.asarray(x), **params)
-                _m = func(getmaskarray(x), **params)
-                return masked_array(_d, mask=_m)
+        if isinstance(x, ndarray):
+            _d = func(x.__array__(), *args, **params)
+            _m = func(getmaskarray(x), *args, **params)
+            return masked_array(_d, mask=_m)
         else:
-            arrays = []
-            args = list(args)
-            while len(args) > 0 and issequence(args[0]):
-                arrays.append(args.pop(0))
-            res = []
-            for x in arrays:
-                _d = func(np.asarray(x), *args, **params)
-                _m = func(getmaskarray(x), *args, **params)
-                res.append(masked_array(_d, mask=_m))
-            return res
-
-atleast_1d = _fromnxfunction('atleast_1d')
-atleast_2d = _fromnxfunction('atleast_2d')
-atleast_3d = _fromnxfunction('atleast_3d')
-#atleast_1d = np.atleast_1d
-#atleast_2d = np.atleast_2d
-#atleast_3d = np.atleast_3d
-
-vstack = row_stack = _fromnxfunction('vstack')
-hstack = _fromnxfunction('hstack')
-column_stack = _fromnxfunction('column_stack')
-dstack = _fromnxfunction('dstack')
-
-hsplit = _fromnxfunction('hsplit')
-
-diagflat = _fromnxfunction('diagflat')
+            _d = func(np.asarray(x), *args, **params)
+            _m = func(getmaskarray(x), *args, **params)
+            return masked_array(_d, mask=_m)
+
+
+class _fromnxfunction_seq(_fromnxfunction):
+    """
+    A version of `_fromnxfunction` that is called with a single sequence
+    of arrays followed by auxiliary args that are passed verbatim for
+    both the data and mask calls.
+    """
+    def __call__(self, x, *args, **params):
+        func = getattr(np, self.__name__)
+        _d = func(tuple([np.asarray(a) for a in x]), *args, **params)
+        _m = func(tuple([getmaskarray(a) for a in x]), *args, **params)
+        return masked_array(_d, mask=_m)
+
+
+class _fromnxfunction_args(_fromnxfunction):
+    """
+    A version of `_fromnxfunction` that is called with multiple array
+    arguments. The first non-array-like input marks the beginning of the
+    arguments that are passed verbatim for both the data and mask calls.
+    Array arguments are processed independently and the results are
+    returned in a list. If only one array is found, the return value is
+    just the processed array instead of a list.
+    """
+    def __call__(self, *args, **params):
+        func = getattr(np, self.__name__)
+        arrays = []
+        args = list(args)
+        while len(args) > 0 and issequence(args[0]):
+            arrays.append(args.pop(0))
+        res = []
+        for x in arrays:
+            _d = func(np.asarray(x), *args, **params)
+            _m = func(getmaskarray(x), *args, **params)
+            res.append(masked_array(_d, mask=_m))
+        if len(arrays) == 1:
+            return res[0]
+        return res
+
+
+class _fromnxfunction_allargs(_fromnxfunction):
+    """
+    A version of `_fromnxfunction` that is called with multiple array
+    arguments. Similar to `_fromnxfunction_args` except that all args
+    are converted to arrays even if they are not so already. This makes
+    it possible to process scalars as 1-D arrays. Only keyword arguments
+    are passed through verbatim for the data and mask calls. Arrays
+    arguments are processed independently and the results are returned
+    in a list. If only one arg is present, the return value is just the
+    processed array instead of a list.
+    """
+    def __call__(self, *args, **params):
+        func = getattr(np, self.__name__)
+        res = []
+        for x in args:
+            _d = func(np.asarray(x), **params)
+            _m = func(getmaskarray(x), **params)
+            res.append(masked_array(_d, mask=_m))
+        if len(args) == 1:
+            return res[0]
+        return res
+
+
+atleast_1d = _fromnxfunction_allargs('atleast_1d')
+atleast_2d = _fromnxfunction_allargs('atleast_2d')
+atleast_3d = _fromnxfunction_allargs('atleast_3d')
+
+vstack = row_stack = _fromnxfunction_seq('vstack')
+hstack = _fromnxfunction_seq('hstack')
+column_stack = _fromnxfunction_seq('column_stack')
+dstack = _fromnxfunction_seq('dstack')
+
+hsplit = _fromnxfunction_single('hsplit')
+
+diagflat = _fromnxfunction_single('diagflat')
 
 
 #####--------------------------------------------------------------------------
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index bb59aad96..33c4b1922 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -1202,7 +1202,7 @@ class TestArraySetOps(TestCase):
 
 class TestShapeBase(TestCase):
 
-    def test_atleast2d(self):
+    def test_atleast_2d(self):
         # Test atleast_2d
         a = masked_array([0, 1, 2], mask=[0, 1, 0])
         b = atleast_2d(a)
@@ -1210,21 +1210,46 @@ class TestShapeBase(TestCase):
         assert_equal(b.mask.shape, b.data.shape)
         assert_equal(a.shape, (3,))
         assert_equal(a.mask.shape, a.data.shape)
+        assert_equal(b.mask.shape, b.data.shape)
 
     def test_shape_scalar(self):
         # the atleast and diagflat function should work with scalars
         # GitHub issue #3367
+        # Additionally, the atleast functions should accept multiple scalars
+        # correctly
         b = atleast_1d(1.0)
-        assert_equal(b.shape, (1, ))
-        assert_equal(b.mask.shape, b.data.shape)
+        assert_equal(b.shape, (1,))
+        assert_equal(b.mask.shape, b.shape)
+        assert_equal(b.data.shape, b.shape)
+
+        b = atleast_1d(1.0, 2.0)
+        for a in b:
+            assert_equal(a.shape, (1,))
+            assert_equal(a.mask.shape, a.shape)
+            assert_equal(a.data.shape, a.shape)
 
         b = atleast_2d(1.0)
         assert_equal(b.shape, (1, 1))
-        assert_equal(b.mask.shape, b.data.shape)
+        assert_equal(b.mask.shape, b.shape)
+        assert_equal(b.data.shape, b.shape)
+
+        b = atleast_2d(1.0, 2.0)
+        for a in b:
+            assert_equal(a.shape, (1, 1))
+            assert_equal(a.mask.shape, a.shape)
+            assert_equal(a.data.shape, a.shape)
 
         b = atleast_3d(1.0)
         assert_equal(b.shape, (1, 1, 1))
-        assert_equal(b.mask.shape, b.data.shape)
+        assert_equal(b.mask.shape, b.shape)
+        assert_equal(b.data.shape, b.shape)
+
+        b = atleast_3d(1.0, 2.0)
+        for a in b:
+            assert_equal(a.shape, (1, 1, 1))
+            assert_equal(a.mask.shape, a.shape)
+            assert_equal(a.data.shape, a.shape)
+
 
         b = diagflat(1.0)
         assert_equal(b.shape, (1, 1))