numpy · madphysicist · Feb 10, 2021 · pbrod · May 9, 2021 · pbrod
diff --git a/doc/release/upcoming_changes/18386.new_function.rst b/doc/release/upcoming_changes/18386.new_function.rst
@@ -0,0 +1,4 @@
+Added `atleast_nd` function to `numpy` and `numpy.ma`
+-----------------------------------------------------
+`atleast_nd` generalizes ``atleast_1d``, ``atleast_2d`` and ``atleast_3d`` to
+arbitrary numbers of dimensions.
diff --git a/doc/source/reference/routines.array-manipulation.rst b/doc/source/reference/routines.array-manipulation.rst
@@ -41,6 +41,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
@@ -127,6 +127,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
@@ -967,6 +967,7 @@ Conversions
   `atleast_1d`,
   `atleast_2d`,
   `atleast_3d`,
+   `atleast_nd`,
   `mat`
 Manipulations
    `array_split`,

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -326,6 +326,7 @@ from numpy.core.shape_base import (
    atleast_1d as atleast_1d,
    atleast_2d as atleast_2d,
    atleast_3d as atleast_3d,
+    atleast_nd as atleast_nd,
    block as block,
    hstack as hstack,
    stack as stack,

diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
@@ -1,5 +1,5 @@
-__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
-           'stack', 'vstack']
+__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'atleast_nd',
+           'block', 'hstack', 'stack', 'vstack']

 import functools
 import itertools
@@ -17,6 +17,12 @@
    overrides.array_function_dispatch, module='numpy')


+def _unpack(lst):
+    if len(lst) == 1:
+        return lst[0]
+    return lst
+
+
 def _atleast_1d_dispatcher(*arys):
    return arys

@@ -42,7 +48,7 @@ def atleast_1d(*arys):

    See Also
    --------
-    atleast_2d, atleast_3d
+    atleast_2d, atleast_3d, atleast_nd

    Examples
    --------
@@ -61,18 +67,7 @@ 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
+    return _unpack([atleast_nd(a, 1, 0) for a in arys])


 def _atleast_2d_dispatcher(*arys):
@@ -87,9 +82,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
    -------
@@ -100,7 +95,7 @@ def atleast_2d(*arys):

    See Also
    --------
-    atleast_1d, atleast_3d
+    atleast_1d, atleast_3d, atleast_nd

    Examples
    --------
@@ -117,20 +112,7 @@ 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[_nx.newaxis, :]
-        else:
-            result = ary
-        res.append(result)
-    if len(res) == 1:
-        return res[0]
-    else:
-        return res
+    return _unpack([atleast_nd(a, 2, 0) for a in arys])


 def _atleast_3d_dispatcher(*arys):
@@ -145,22 +127,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
+    function are not consistent 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
    --------
@@ -187,22 +178,105 @@ def atleast_3d(*arys):
    [[[1 2]]] (1, 1, 2)

    """
-    res = []
-    for ary in arys:
-        ary = asanyarray(ary)
-        if ary.ndim == 0:
-            result = ary.reshape(1, 1, 1)
-        elif ary.ndim == 1:
-            result = ary[_nx.newaxis, :, _nx.newaxis]
-        elif ary.ndim == 2:
-            result = ary[:, :, _nx.newaxis]
-        else:
-            result = ary
-        res.append(result)
-    if len(res) == 1:
-        return res[0]
-    else:
-        return res
+    return _unpack([atleast_nd(atleast_nd(a, 2, 0), 3, -1) for a in arys])
+
+
+def _atleast_nd_dispatcher(ary, ndim, pos=None):
+    return (ary,)
+
+
+@array_function_dispatch(_atleast_nd_dispatcher)
+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 : int
+        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, 2) for x in ((1, 2), [[3, 4]], [[[5, 6]]])]
+    [array([[1, 2]]), array([[3, 4]]), array([[[5, 6]]])]
+
+    >>> 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)
+    pos = normalize_axis_index(pos, ary.ndim + 1)
+    extra = operator.index(ndim) - ary.ndim
+    if extra > 0:
+        ind = pos * (slice(None),) + extra * (None,) + (Ellipsis,)
+        ary = ary[ind]
+    return ary


 def _arrays_for_stack_dispatcher(arrays, stacklevel=4):

diff --git a/numpy/core/shape_base.pyi b/numpy/core/shape_base.pyi
@@ -26,6 +26,8 @@ def atleast_3d(__arys: ArrayLike) -> ndarray: ...
 @overload
 def atleast_3d(*arys: ArrayLike) -> List[ndarray]: ...

+def atleast_nd(ary: ArrayLike, ndim: int, pos: int = ...) -> ndarray: ...
+
 def vstack(tup: Sequence[ArrayLike]) -> ndarray: ...
 def hstack(tup: Sequence[ArrayLike]) -> ndarray: ...
 @overload

diff --git a/numpy/core/tests/test_shape_base.py b/numpy/core/tests/test_shape_base.py
@@ -1,9 +1,9 @@
+from decimal import Decimal
 import pytest
 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, vstack, hstack, newaxis, concatenate, stack)
 from numpy.core.shape_base import (_block_dispatcher, _block_setup,
                                   _block_concatenate, _block_slicing)
 from numpy.testing import (
@@ -123,6 +123,71 @@ def test_3D_array(self):
        assert_array_equal(res, desired)


+class TestAtleastNd(object):
+    def test_0D_arrays(self):
+        a = array(3)
+        dims = [3, 2, 0]
+        expected = [array([[[3]]]), array([[3]]), array(3)]
+
+        for b, d in zip(expected, dims):
+            assert_array_equal(atleast_nd(a, d), b)
+            assert_array_equal(atleast_nd(a, d, -1), b)
+
+    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)
+
+    def test_ndim(self):
+        a = 3
+        assert_raises(TypeError, atleast_nd, a, 0.4)
+        assert_raises(TypeError, atleast_nd, a, Decimal(4))
+        assert_raises(TypeError, atleast_nd, a, np.array([4, 5]))
+        assert_raises(np.AxisError, atleast_nd, a, -2, 1)
+        assert_equal(atleast_nd(a, np.array(4, dtype=np.uint8)).ndim, 4)
+
+        assert isinstance(atleast_nd(a, 0, 0), np.ndarray)
+        assert_equal(atleast_nd(a, -5).ndim, 0)
+        assert_equal(atleast_nd(a, -5, -1).ndim, 0)
+
+
 class TestHstack:
    def test_non_iterable(self):
        assert_raises(TypeError, hstack, 1)

diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
@@ -542,7 +542,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
    --------