Author: Brian Kearns <bdkea...@gmail.com>
Branch: 
Changeset: r61085:5fb4c2227097
Date: 2013-02-11 06:31 -0500
http://bitbucket.org/pypy/pypy/changeset/5fb4c2227097/

Log:    might as well add dstack too

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
@@ -272,3 +272,52 @@
     else:
         return _numpypy.concatenate(arrs, 1)
 
+def dstack(tup):
+    """
+    Stack arrays in sequence depth wise (along third axis).
+
+    Takes a sequence of arrays and stack them along the third axis
+    to make a single array. Rebuilds arrays divided by `dsplit`.
+    This is a simple way to stack 2D arrays (images) into a single
+    3D array for processing.
+
+    Parameters
+    ----------
+    tup : sequence of arrays
+        Arrays to stack. All of them must have the same shape along all
+        but the third axis.
+
+    Returns
+    -------
+    stacked : ndarray
+        The array formed by stacking the given arrays.
+
+    See Also
+    --------
+    vstack : Stack along first axis.
+    hstack : Stack along second axis.
+    concatenate : Join arrays.
+    dsplit : Split array along third axis.
+
+    Notes
+    -----
+    Equivalent to ``np.concatenate(tup, axis=2)``.
+
+    Examples
+    --------
+    >>> a = np.array((1,2,3))
+    >>> b = np.array((2,3,4))
+    >>> np.dstack((a,b))
+    array([[[1, 2],
+            [2, 3],
+            [3, 4]]])
+
+    >>> a = np.array([[1],[2],[3]])
+    >>> b = np.array([[2],[3],[4]])
+    >>> np.dstack((a,b))
+    array([[[1, 2]],
+           [[2, 3]],
+           [[3, 4]]])
+
+    """
+    return _numpypy.concatenate(map(atleast_3d,tup),2)
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
@@ -94,3 +94,14 @@
                                   [2, 3],
                                   [3, 4]])
 
+    def test_dstack(self):
+        import numpypy as np
+        a = np.array((1,2,3))
+        b = np.array((2,3,4))
+        c = np.dstack((a,b))
+        assert np.array_equal(c, [[[1, 2], [2, 3], [3, 4]]])
+
+        a = np.array([[1],[2],[3]])
+        b = np.array([[2],[3],[4]])
+        c = np.dstack((a,b))
+        assert np.array_equal(c, [[[1, 2]], [[2, 3]], [[3, 4]]])
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
http://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to