[pypy-commit] pypy default: implement numpypy.atleast_3d (by copying it from numpy source)

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: 
Changeset: r61082:041283a61f06
Date: 2013-02-11 11:43 +0100
http://bitbucket.org/pypy/pypy/changeset/041283a61f06/

Log:	implement numpypy.atleast_3d (by copying it from numpy source)

diff --git a/lib_pypy/numpypy/core/shape_base.py b/lib_pypy/numpypy/core/shape_base.py
--- a/lib_pypy/numpypy/core/shape_base.py
+++ b/lib_pypy/numpypy/core/shape_base.py
@@ -104,3 +104,69 @@
         return res[0]
     else:
         return res
+
+def atleast_3d(*arys):
+    """
+    View inputs as arrays with at least three dimensions.
+
+    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.
+
+    Returns
+    -------
+    res1, res2, ... : ndarray
+        An array, or tuple 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
+
+    Examples
+    --------
+    >>> np.atleast_3d(3.0)
+    array([[[ 3.]]])
+
+    >>> x = np.arange(3.0)
+    >>> np.atleast_3d(x).shape
+    (1, 3, 1)
+
+    >>> x = np.arange(12.0).reshape(4,3)
+    >>> np.atleast_3d(x).shape
+    (4, 3, 1)
+    >>> np.atleast_3d(x).base is x
+    True
+
+    >>> for arr in np.atleast_3d([1, 2], [[1, 2]], [[[1, 2]]]):
+    ...     print arr, arr.shape
+    ...
+    [[[1]
+      [2]]] (1, 2, 1)
+    [[[1]
+      [2]]] (1, 2, 1)
+    [[[1 2]]] (1, 1, 2)
+
+    """
+    res = []
+    for ary in arys:
+        ary = asanyarray(ary)
+        if len(ary.shape) == 0:
+            result = ary.reshape(1,1,1)
+        elif len(ary.shape) == 1:
+            result = ary[newaxis,:,newaxis]
+        elif len(ary.shape) == 2:
+            result = ary[:,:,newaxis]
+        else:
+            result = ary
+        res.append(result)
+    if len(res) == 1:
+        return res[0]
+    else:
+        return res
diff --git a/pypy/module/test_lib_pypy/numpypy/core/test_shape_base.py b/pypy/module/test_lib_pypy/numpypy/core/test_shape_base.py
--- a/pypy/module/test_lib_pypy/numpypy/core/test_shape_base.py
+++ b/pypy/module/test_lib_pypy/numpypy/core/test_shape_base.py
@@ -34,3 +34,29 @@
         assert np.array_equal(a[0], [[1]])
         assert np.array_equal(a[1], [[1, 2]])
         assert np.array_equal(a[2], [[1, 2]])
+
+    def test_atleast_3d(self):
+        import numpypy as np
+
+        a = np.atleast_3d(3.0)
+        assert np.array_equal(a, [[[ 3.]]])
+
+        x = np.arange(3.0)
+        assert np.atleast_3d(x).shape == (1, 3, 1)
+
+        x = np.arange(12.0).reshape(4,3)
+        assert np.atleast_3d(x).shape == (4, 3, 1)
+
+        a = np.atleast_3d([1, 2])
+        assert np.array_equal(a, [[[1],
+                                   [2]]])
+        assert a.shape == (1, 2, 1)
+
+        a = np.atleast_3d([[1, 2]])
+        assert np.array_equal(a, [[[1],
+                                   [2]]])
+        assert a.shape == (1, 2, 1)
+
+        a = np.atleast_3d([[[1, 2]]])
+        assert np.array_equal(a, [[[1, 2]]])
+        assert a.shape == (1, 1, 2)