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

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: 
Changeset: r61084:74a2914159df
Date: 2013-02-11 11:52 +0100
http://bitbucket.org/pypy/pypy/changeset/74a2914159df/

Log:	implement numpypy.hstack (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
@@ -222,3 +222,53 @@
 
     """
     return _numpypy.concatenate(map(atleast_2d,tup),0)
+
+def hstack(tup):
+    """
+    Stack arrays in sequence horizontally (column wise).
+
+    Take a sequence of arrays and stack them horizontally to make
+    a single array. Rebuild arrays divided by `hsplit`.
+
+    Parameters
+    ----------
+    tup : sequence of ndarrays
+        All arrays must have the same shape along all but the second axis.
+
+    Returns
+    -------
+    stacked : ndarray
+        The array formed by stacking the given arrays.
+
+    See Also
+    --------
+    vstack : Stack arrays in sequence vertically (row wise).
+    dstack : Stack arrays in sequence depth wise (along third axis).
+    concatenate : Join a sequence of arrays together.
+    hsplit : Split array along second axis.
+
+    Notes
+    -----
+    Equivalent to ``np.concatenate(tup, axis=1)``
+
+    Examples
+    --------
+    >>> a = np.array((1,2,3))
+    >>> b = np.array((2,3,4))
+    >>> np.hstack((a,b))
+    array([1, 2, 3, 2, 3, 4])
+    >>> a = np.array([[1],[2],[3]])
+    >>> b = np.array([[2],[3],[4]])
+    >>> np.hstack((a,b))
+    array([[1, 2],
+           [2, 3],
+           [3, 4]])
+
+    """
+    arrs = map(atleast_1d,tup)
+    # As a special case, dimension 0 of 1-dimensional arrays is "horizontal"
+    if arrs[0].ndim == 1:
+        return _numpypy.concatenate(arrs, 0)
+    else:
+        return _numpypy.concatenate(arrs, 1)
+
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
@@ -79,3 +79,18 @@
                                   [2],
                                   [3],
                                   [4]])
+
+    def test_hstack(self):
+        import numpypy as np
+        a = np.array((1,2,3))
+        b = np.array((2,3,4))
+        c = np.hstack((a,b))
+        assert np.array_equal(c, [1, 2, 3, 2, 3, 4])
+        
+        a = np.array([[1],[2],[3]])
+        b = np.array([[2],[3],[4]])
+        c = np.hstack((a,b))
+        assert np.array_equal(c, [[1, 2],
+                                  [2, 3],
+                                  [3, 4]])
+