diff --git a/doc/release/1.16.0-notes.rst b/doc/release/1.16.0-notes.rst index 2b84bb90a311..a2098e040786 100644 --- a/doc/release/1.16.0-notes.rst +++ b/doc/release/1.16.0-notes.rst @@ -19,6 +19,9 @@ Highlights New functions ============= +* `atleast_nd` generalizes ``atleast_1d``, ``atleast_2d`` and ``atleast_3d`` to + arbitrary numbers of dimensions. + Deprecations ============ diff --git a/doc/source/reference/routines.array-manipulation.rst b/doc/source/reference/routines.array-manipulation.rst index cc93d1029f52..88928c551325 100644 --- a/doc/source/reference/routines.array-manipulation.rst +++ b/doc/source/reference/routines.array-manipulation.rst @@ -40,6 +40,7 @@ Changing number of dimensions atleast_1d atleast_2d atleast_3d + atleast_nd broadcast broadcast_to broadcast_arrays diff --git a/doc/source/reference/routines.ma.rst b/doc/source/reference/routines.ma.rst index 15f2ba0a4add..2cbe6f5d86cb 100644 --- a/doc/source/reference/routines.ma.rst +++ b/doc/source/reference/routines.ma.rst @@ -121,6 +121,7 @@ Changing the number of dimensions ma.atleast_1d ma.atleast_2d ma.atleast_3d + ma.atleast_nd ma.expand_dims ma.squeeze diff --git a/doc/source/user/quickstart.rst b/doc/source/user/quickstart.rst index 5ef8b145f6be..8337fd8896be 100644 --- a/doc/source/user/quickstart.rst +++ b/doc/source/user/quickstart.rst @@ -914,6 +914,7 @@ Conversions `atleast_1d`, `atleast_2d`, `atleast_3d`, + `atleast_nd`, `mat` Manipulations `array_split`, diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py index 30919ed7e960..1092db4df2a9 100644 --- a/numpy/core/shape_base.py +++ b/numpy/core/shape_base.py @@ -1,7 +1,7 @@ from __future__ import division, absolute_import, print_function -__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack', - 'stack', 'vstack'] +__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'atleast_nd', 'block', + 'hstack', 'stack', 'vstack'] from . import numeric as _nx @@ -9,6 +9,7 @@ from .multiarray import normalize_axis_index from ._internal import recursive + def atleast_1d(*arys): """ Convert inputs to arrays with at least one dimension. @@ -29,7 +30,7 @@ def atleast_1d(*arys): See Also -------- - atleast_2d, atleast_3d + atleast_2d, atleast_3d, atleast_nd Examples -------- @@ -48,18 +49,10 @@ def atleast_1d(*arys): [array([1]), array([3, 4])] """ - res = [] - for ary in arys: - ary = asanyarray(ary) - if ary.ndim == 0: - result = ary.reshape(1) - else: - result = ary - res.append(result) - if len(res) == 1: - return res[0] - else: - return res + if len(arys) == 1: + return atleast_nd(arys[0], 1) + return [atleast_nd(x, 1) for x in arys] + def atleast_2d(*arys): """ @@ -68,9 +61,9 @@ def atleast_2d(*arys): Parameters ---------- arys1, arys2, ... : array_like - One or more array-like sequences. Non-array inputs are converted - to arrays. Arrays that already have two or more dimensions are - preserved. + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have two or more + dimensions are preserved. Returns ------- @@ -81,7 +74,7 @@ def atleast_2d(*arys): See Also -------- - atleast_1d, atleast_3d + atleast_1d, atleast_3d, atleast_nd Examples -------- @@ -98,20 +91,10 @@ def atleast_2d(*arys): [array([[1]]), array([[1, 2]]), array([[1, 2]])] """ - res = [] - for ary in arys: - ary = asanyarray(ary) - if ary.ndim == 0: - result = ary.reshape(1, 1) - elif ary.ndim == 1: - result = ary[newaxis,:] - else: - result = ary - res.append(result) - if len(res) == 1: - return res[0] - else: - return res + if len(arys) == 1: + return atleast_nd(arys[0], 2) + return [atleast_nd(x, 2) for x in arys] + def atleast_3d(*arys): """ @@ -120,22 +103,31 @@ def atleast_3d(*arys): Parameters ---------- arys1, arys2, ... : array_like - One or more array-like sequences. Non-array inputs are converted to - arrays. Arrays that already have three or more dimensions are - preserved. + One or more array-like sequences. Non-array inputs are + converted to arrays. Arrays that already have three or more + dimensions are preserved. Returns ------- res1, res2, ... : ndarray - An array, or list of arrays, each with ``a.ndim >= 3``. Copies are - avoided where possible, and views with three or more dimensions are - returned. For example, a 1-D array of shape ``(N,)`` becomes a view - of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a - view of shape ``(M, N, 1)``. + An array, or list of arrays, each with ``a.ndim >= 3``. Copies + are avoided where possible, and views with three or more + dimensions are returned. For example, a 1-D array of shape + ``(N,)`` becomes a view of shape ``(1, N, 1)``, and a 2-D array + of shape ``(M, N)`` becomes a view of shape ``(M, N, 1)``. See Also -------- - atleast_1d, atleast_2d + atleast_1d, atleast_2d, atleast_nd + + Notes + ----- + As mentioned in the `Returns` section, the results of this fuction + are not compatible with any of the other `atleast*` functions. + `atleast_2d` prepends the unit dimension to a 1D array while + `atleast_3d` appends it to a 2D array. The 1D array case both + appends and prepends a dimension, while `atleast_nd` can only add + dimensions to one end at a time. Examples -------- @@ -168,9 +160,9 @@ def atleast_3d(*arys): if ary.ndim == 0: result = ary.reshape(1, 1, 1) elif ary.ndim == 1: - result = ary[newaxis,:, newaxis] + result = ary[newaxis, :, newaxis] elif ary.ndim == 2: - result = ary[:,:, newaxis] + result = ary[:, :, newaxis] else: result = ary res.append(result) @@ -180,6 +172,100 @@ def atleast_3d(*arys): return res +def atleast_nd(ary, ndim, pos=0): + """ + View input as array with at least `ndim` dimensions. + + New unit dimensions are inserted at the index given by `pos` if + necessary. + + Parameters + ---------- + ary : array_like + The input array. Non-array inputs are converted to arrays. + Arrays that already have `ndim` or more dimensions are + preserved. + ndim : scalar + The minimum number of dimensions required. + pos : int, optional + The index to insert the new dimensions. May range from + ``-ary.ndim - 1`` to ``+ary.ndim`` (inclusive). Non-negative + indices indicate locations before the corresponding axis: + ``pos=0`` means to insert at the very beginning. Negative + indices indicate locations after the corresponding axis: + ``pos=-1`` means to insert at the very end. 0 and -1 are always + guaranteed to work. Any other number will depend on the + dimensions of the existing array. Default is 0. + + Returns + ------- + res : ndarray + An array with ``res.ndim >= ndim``. A view is returned for array + inputs. Dimensions are prepended if `pos` is 0, so for example, + a 1-D array of shape ``(N,)`` with ``ndim=4`` becomes a view of + shape ``(1, 1, 1, N)``. Dimensions are appended if `pos` is -1, + so for example a 2-D array of shape ``(M, N)`` becomes a view of + shape ``(M, N, 1, 1)`` when ``ndim=4``. + + See Also + -------- + atleast_1d, atleast_2d, atleast_3d + + Notes + ----- + This function does not follow the convention of the other atleast_*d + functions in numpy in that it only accepts a single array argument. + To process multiple arrays, use a comprehension or loop around the + function call. See examples below. + + Setting ``pos=0`` is equivalent to how the array would be + interpreted by numpy's broadcasting rules. There is no need to call + this function for simple broadcasting. This is also roughly + (but not exactly) equivalent to + ``np.array(ary, copy=False, subok=True, ndmin=ndim)``. + + It is easy to create functions for specific dimensions similar to + the other atleast_*d functions using Python's `functools.partial` + function. An example is shown below. + + Examples + -------- + >>> np.atleast_nd(3.0, 4) + array([[[[ 3.]]]]) + + >>> x = np.arange(3.0) + >>> np.atleast_nd(x, 2).shape + (1, 3) + + >>> x = np.arange(12.0).reshape(4, 3) + >>> np.atleast_nd(x, 5).shape + (1, 1, 1, 4, 3) + >>> np.atleast_nd(x, 5).base is x.base + True + + >>> [np.atleast_nd(x) for x in ((1, 2), [[1, 2]], [[[1, 2]]])]: + [array([[1, 2]]), array([[1, 2]]), array([[[1, 2]]])] + + >>> np.atleast_nd((1, 2), 5, pos=0).shape + (1, 1, 1, 1, 2) + >>> np.atleast_nd((1, 2), 5, pos=-1).shape + (2, 1, 1, 1, 1) + + >>> from functools import partial + >>> atleast_4d = partial(np.atleast_nd, ndim=4) + >>> atleast_4d([1, 2, 3]) + [[[[1, 2, 3]]]] + """ + ary = array(ary, copy=False, subok=True) + if ary.ndim: + pos = normalize_axis_index(pos, ary.ndim + 1) + extra = ndim - ary.ndim + if extra > 0: + ind = pos * (slice(None),) + extra * (None,) + (Ellipsis,) + ary = ary[ind] + return ary + + def vstack(tup): """ Stack arrays in sequence vertically (row wise). diff --git a/numpy/core/tests/test_shape_base.py b/numpy/core/tests/test_shape_base.py index 72b3451a4f8c..474b99d79c8b 100644 --- a/numpy/core/tests/test_shape_base.py +++ b/numpy/core/tests/test_shape_base.py @@ -3,8 +3,8 @@ import warnings import numpy as np from numpy.core import ( - array, arange, atleast_1d, atleast_2d, atleast_3d, block, vstack, hstack, - newaxis, concatenate, stack + array, arange, atleast_1d, atleast_2d, atleast_3d, atleast_nd, block, + concatenate, hstack, newaxis, stack, vstack, ) from numpy.testing import ( assert_, assert_raises, assert_array_equal, assert_equal, @@ -125,6 +125,59 @@ def test_3D_array(self): assert_array_equal(res, desired) +class TestAtleastNd(object): + def test_0D_arrays(self): + arrays = [array(1), array(2)] + dims = [3, 2] + expected = [array([[[1]]]), array([[2]])] + + for x, y, d in zip(arrays, expected, dims): + assert_array_equal(atleast_nd(x, d), y) + assert_array_equal(atleast_nd(x, d, -1), y) + + def test_nD_arrays(self): + a = array([1]) + b = array([4, 5, 6]) + c = array([[2, 3]]) + d = array([[[2], [3]], [[2], [3]]]) + e = ((((1, 2), (3, 4)), ((5, 6), (7, 8)))) + arrays = (a, b, c, d, e) + expected_before = (array([[[1]]]), + array([[[4, 5, 6]]]), + array([[[2, 3]]]), + d, + array(e)) + expected_after = (array([[[1]]]), + array([[[4]], [[5]], [[6]]]), + array([[[2], [3]]]), + d, + array(e)) + + for x, y in zip(arrays, expected_before): + assert_array_equal(atleast_nd(x, 3), y) + for x, y in zip(arrays, expected_after): + assert_array_equal(atleast_nd(x, 3, pos=-1), y) + + def test_nocopy(self): + a = arange(12.0).reshape(4, 3) + res = atleast_nd(a, 5) + desired_shape = (1, 1, 1, 4, 3) + desired_base = a.base # a was reshaped + assert_equal(res.shape, desired_shape) + assert_(res.base is desired_base) + + def test_passthough(self): + a = array([1, 2, 3]) + assert_(atleast_nd(a, 0) is a) + assert_(atleast_nd(a, 1) is a) + + def test_other_pos(self): + a = arange(12.0).reshape(4, 3) + res = atleast_nd(a, 4, pos=1) + assert_equal(res.shape, (4, 1, 1, 3)) + assert_raises(ValueError, atleast_nd, a, 4, pos=5) + + class TestHstack(object): def test_non_iterable(self): assert_raises(TypeError, hstack, 1) diff --git a/numpy/lib/info.py b/numpy/lib/info.py index 8815a52f0a64..ca5ab8fac4fe 100644 --- a/numpy/lib/info.py +++ b/numpy/lib/info.py @@ -70,6 +70,7 @@ atleast_1d Force arrays to be >= 1D atleast_2d Force arrays to be >= 2D atleast_3d Force arrays to be >= 3D +atleast_nd Force arrays to be >= ND, where N is specified 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/lib/shape_base.py b/numpy/lib/shape_base.py index 66f53473485d..c7240f729267 100644 --- a/numpy/lib/shape_base.py +++ b/numpy/lib/shape_base.py @@ -506,7 +506,7 @@ def expand_dims(a, axis): -------- squeeze : The inverse operation, removing singleton dimensions reshape : Insert, remove, and combine dimensions, and resize existing ones - doc.indexing, atleast_1d, atleast_2d, atleast_3d + doc.indexing, atleast_1d, atleast_2d, atleast_3d, atleast_nd Examples -------- diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py index 3be4d3625b58..93d69366921e 100644 --- a/numpy/ma/extras.py +++ b/numpy/ma/extras.py @@ -12,7 +12,7 @@ __all__ = [ 'apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d', - 'atleast_3d', 'average', 'clump_masked', 'clump_unmasked', + 'atleast_3d', 'atleast_nd', 'average', 'clump_masked', 'clump_unmasked', 'column_stack', 'compress_cols', 'compress_nd', 'compress_rowcols', 'compress_rows', 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot', 'dstack', 'ediff1d', 'flatnotmasked_contiguous', 'flatnotmasked_edges', @@ -280,12 +280,10 @@ def __call__(self, x, *args, **params): func = getattr(np, self.__name__) if isinstance(x, ndarray): _d = func(x.__array__(), *args, **params) - _m = func(getmaskarray(x), *args, **params) - return masked_array(_d, mask=_m) else: _d = func(np.asarray(x), *args, **params) - _m = func(getmaskarray(x), *args, **params) - return masked_array(_d, mask=_m) + _m = func(getmaskarray(x), *args, **params) + return masked_array(_d, mask=_m) class _fromnxfunction_seq(_fromnxfunction): @@ -352,6 +350,7 @@ def __call__(self, *args, **params): atleast_1d = _fromnxfunction_allargs('atleast_1d') atleast_2d = _fromnxfunction_allargs('atleast_2d') atleast_3d = _fromnxfunction_allargs('atleast_3d') +atleast_nd = _fromnxfunction_single('atleast_nd') vstack = row_stack = _fromnxfunction_seq('vstack') hstack = _fromnxfunction_seq('hstack') diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py index c29bec2bdf84..85efcc0f3290 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -24,12 +24,12 @@ nomask, ones, zeros, count ) from numpy.ma.extras import ( - atleast_1d, atleast_2d, atleast_3d, mr_, dot, polyfit, cov, corrcoef, - median, average, unique, setxor1d, setdiff1d, union1d, intersect1d, in1d, - ediff1d, apply_over_axes, apply_along_axis, compress_nd, compress_rowcols, - mask_rowcols, clump_masked, clump_unmasked, flatnotmasked_contiguous, - notmasked_contiguous, notmasked_edges, masked_all, masked_all_like, isin, - diagflat, stack, vstack, hstack + apply_along_axis, apply_over_axes, atleast_1d, atleast_2d, atleast_3d, + atleast_nd, average, clump_masked, clump_unmasked, compress_nd, + compress_rowcols, corrcoef, cov, diagflat, dot, ediff1d, + flatnotmasked_contiguous, in1d, intersect1d, isin, mask_rowcols, masked_all, + masked_all_like, median, mr_, notmasked_contiguous, notmasked_edges, + polyfit, setxor1d, setdiff1d, stack, union1d, unique, vstack ) import numpy.ma.extras as mae @@ -1553,6 +1553,25 @@ def test_atleast_2d(self): assert_equal(a.mask.shape, a.data.shape) assert_equal(b.mask.shape, b.data.shape) + def test_atleast_nd(self): + # Test atleast_nd + a = masked_array([0, 1, 2], mask=[0, 1, 0]) + b = atleast_nd(a, 4) + c = atleast_nd(a, 4, pos=-1) + d = atleast_nd(a, 4, pos=1) + assert_equal(a.shape, (3,)) + assert_equal(a.data.shape, a.shape) + assert_equal(a.mask.shape, a.shape) + assert_equal(b.shape, (1, 1, 1, 3)) + assert_equal(b.data.shape, b.shape) + assert_equal(b.mask.shape, b.shape) + assert_equal(c.shape, (3, 1, 1, 1)) + assert_equal(c.data.shape, c.shape) + assert_equal(c.mask.shape, c.shape) + assert_equal(d.shape, (3, 1, 1, 1)) + assert_equal(d.data.shape, d.shape) + assert_equal(d.mask.shape, d.shape) + def test_shape_scalar(self): # the atleast and diagflat function should work with scalars # GitHub issue #3367 @@ -1591,6 +1610,12 @@ def test_shape_scalar(self): assert_equal(a.mask.shape, a.shape) assert_equal(a.data.shape, a.shape) + b = atleast_nd(1.0, 4) + c = atleast_nd(1.0, 4, -1) + assert_equal(b.shape, (1, 1, 1, 1)) + assert_equal(b.mask.shape, b.data.shape) + assert_equal(c.shape, (1, 1, 1, 1)) + assert_equal(c.mask.shape, c.data.shape) b = diagflat(1.0) assert_equal(b.shape, (1, 1))