[pypy-commit] pypy numpy-concatenate: Adds numpy.concatenate

[jterrace]

Author: Jeff Terrace <jterrace at gmail.com>
Branch: numpy-concatenate
Changeset: r50522:0d81f0701f18
Date: 2011-12-14 16:17 -0500
http://bitbucket.org/pypy/pypy/changeset/0d81f0701f18/

Log:	Adds numpy.concatenate

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -97,4 +97,5 @@
         'pi': 'app_numpy.pi',
         'arange': 'app_numpy.arange',
         'reshape': 'app_numpy.reshape',
+        'concatenate': 'app_numpy.concatenate',
     }
diff --git a/pypy/module/micronumpy/app_numpy.py b/pypy/module/micronumpy/app_numpy.py
--- a/pypy/module/micronumpy/app_numpy.py
+++ b/pypy/module/micronumpy/app_numpy.py
@@ -34,6 +34,37 @@
         a = numpypy.array(a)
     return a.max()
 
+def concatenate(array_iter, axis=0):
+    arrays = []
+    shape = None
+    for a in array_iter:
+        if not hasattr(a, 'shape'):
+            a = numpypy.array(a)
+        arrays.append(a)
+        if len(a.shape) < axis + 1:
+            raise ValueError("bad axis argument")
+        if shape is None:
+            shape = list(a.shape)
+        else:
+            for i, axis_size in enumerate(a.shape):
+                if len(a.shape) != len(shape) or (i != axis and axis_size != shape[i]):
+                    raise ValueError("array dimensions must agree except for axis being concatenated")
+                elif i == axis:
+                    shape[i] += axis_size
+    
+    if len(arrays) == 0:
+        raise ValueError("concatenation of zero-length sequences is impossible")
+    
+    out_array = numpypy.zeros(shape)
+    slicing_index = [slice(None)] * len(shape)
+    axis_ptr = 0
+    for a in arrays:
+        slicing_index[axis] = slice(axis_ptr, axis_ptr + a.shape[axis])
+        out_array[tuple(slicing_index)] = a
+        axis_ptr += a.shape[axis]
+    
+    return out_array
+    
 def arange(start, stop=None, step=1, dtype=None):
     '''arange([start], stop[, step], dtype=None)
     Generate values in the half-interval [start, stop).
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -923,6 +923,31 @@
         assert a[:,0].tolist() == [17.1, 40.3]
         assert a[0].tolist() == [17.1, 27.2]
 
+    def test_concatenate(self):
+        from numpypy import array, concatenate
+        a1 = array([0,1,2])
+        a2 = array([3,4,5])
+        a = concatenate((a1, a2))
+        assert len(a) == 6
+        assert (a == [0,1,2,3,4,5]).all()
+        b1 = array([[1, 2], [3, 4]])
+        b2 = array([[5, 6]])
+        b = concatenate((b1, b2), axis=0)
+        assert (b == [[1, 2],[3, 4],[5, 6]]).all()
+        c = concatenate((b1, b2.T), axis=1)
+        assert (c == [[1, 2, 5],[3, 4, 6]]).all()
+        
+        bad_axis = raises(ValueError, concatenate, (a1,a2), axis=1)
+        assert str(bad_axis.value) == "bad axis argument"
+        
+        concat_zero = raises(ValueError, concatenate, ())
+        assert str(concat_zero.value) == \
+            "concatenation of zero-length sequences is impossible"
+        
+        dims_disagree = raises(ValueError, concatenate, (a1, b1), axis=0)
+        assert str(dims_disagree.value) == \
+            "array dimensions must agree except for axis being concatenated"
+
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):