diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py index 31d7d1c0322e..2b5110473bcb 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 ca1e7239756e..141df2ace8aa 100644 --- a/numpy/lib/info.py +++ b/numpy/lib/info.py @@ -67,9 +67,9 @@ ------------------ ================ =================== 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 cbf7b6cdb63b..aefdd4e49a47 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 @@ def getdoc(self): 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 bb59aad96084..33c4b19225e8 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -1202,7 +1202,7 @@ def test_setdiff1d_char_array(self): 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 @@ def test_atleast2d(self): 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))