numpy · madphysicist · Jul 6, 2016 · ahaldane · Oct 2, 2018 · ahaldane
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -1,14 +1,15 @@
 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
 from .numeric import array, asanyarray, newaxis
 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):
    <
8000
span class='blob-code-inner blob-code-marker ' data-code-marker=" ">    [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
@@ -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
@@ -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
@@ -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
     --------