[pypy-commit] pypy separate-applevel-numpy: implement numpy.bincount with tests

[timo_jbo]

Author: Timo Paulssen <timonator at perpetuum-immobile.de>
Branch: separate-applevel-numpy
Changeset: r47638:88a4e8b8d916
Date: 2011-09-27 21:43 +0200
http://bitbucket.org/pypy/pypy/changeset/88a4e8b8d916/

Log:	implement numpy.bincount with tests

diff --git a/lib_pypy/numpy/__init__.py b/lib_pypy/numpy/__init__.py
--- a/lib_pypy/numpy/__init__.py
+++ b/lib_pypy/numpy/__init__.py
@@ -46,3 +46,44 @@
     if not hasattr(a, "mean"):
         a = array(a)
     return a.mean()
+
+__SIGNEDLTR = "i"
+__UNSIGNEDLTR = "u"
+
+def bincount(x, weights=None, minlength=None):
+    if minlength is not None:
+        result = [0 for _ in range(minlength)]
+    else:
+        result = [0]
+
+    if len(x) == 0:
+        raise ValueError("the first argument cannot be empty.")
+
+    x = array(x)
+    if x.dtype.kind not in (__SIGNEDLTR, __UNSIGNEDLTR):
+        raise TypeError("array cannot be safely cast to required type")
+
+
+    if len(x.shape) > 1:
+        raise ValueError("object too deep for desired array")
+
+    if weights is not None:
+        weights = array(weights)
+        if weights.shape != x.shape:
+            raise ValueError("The weights and list don't have the same length.")
+
+        num_iter = (num_and_weight for num_and_weight in zip(x, weights))
+
+    else:
+        num_iter = ((num, 1) for num in x)
+
+    for number, weight in num_iter:
+        if number < 0:
+            raise ValueError("The first argument of bincount must be non-negative")
+        try:
+            result[number] += weight
+        except IndexError:
+            result += [0] * (number - len(result)) + [weight]
+
+    return array(result)
+
diff --git a/lib_pypy/pypy_test/test_numpy.py b/lib_pypy/pypy_test/test_numpy.py
--- a/lib_pypy/pypy_test/test_numpy.py
+++ b/lib_pypy/pypy_test/test_numpy.py
@@ -13,3 +13,64 @@
         from numpy import array, average
         assert average(range(10)) == 4.5
         assert average(array(range(10))) == 4.5
+
+    def w_array_compare(self, have, want):
+        assert len(have) == len(want)
+        mismatch = []
+        for num, (x, y) in enumerate(zip(have, want)):
+            if not x == y:
+                mismatch.append(num)
+        if mismatch:
+            print have
+            print want
+            print mismatch
+        assert mismatch == []
+
+    def test_bincount_simple(self):
+        from numpy import array, bincount
+        a = array(range(10))
+        have, want = (bincount(a), array([1] * 10))
+        self.array_compare(have, want)
+        b = array([9, 9, 9, 9, 9])
+        have, want = (bincount(b), array([0] * 9 + [5]))
+        self.array_compare(have, want)
+        c = [9, 9, 9, 9, 9]
+        have, want = (bincount(c), array([0] * 9 + [5]))
+        self.array_compare(have, want)
+
+    def test_bincount_weights(self):
+        from numpy import array, bincount
+        a = array([1, 2, 3, 3, 4, 5])
+        wa = array([-1, 0.1, 0.2, 0.3, 1, 1])
+        have, want = (bincount(a, wa), array([0, -1, 0.1, 0.5, 1, 1]))
+        self.array_compare(have, want)
+        b = [1, 1, 4]
+        wb = [99, -9, 0]
+        have, want = (bincount(b, wb), array([0, 90, 0, 0, 0]))
+        self.array_compare(have, want)
+        c = [1, 1, 1]
+        wc = [9, 8, 7]
+        have, want = (bincount(c, wc, 4), array([0, 24, 0, 0]))
+        self.array_compare(have, want)
+
+    def test_array_compare(self):
+        # just some sanity-checks for the comparison function
+        from numpy import array
+        a = array([1, 2, 4])
+        b = array([1, 1, 1])
+        raises(AssertionError, "self.array_compare(a, b)")
+        x = array([1, 1, 4])
+        y = array([1, 1])
+        raises(AssertionError, "self.array_compare(x, y)")
+
+    def test_bincount_error_handling(self):
+        from numpy import array, bincount
+        a = array([1.0, 2.0, 3.0], float)
+        raises(TypeError, "bincount(a)")
+        b = array([-1, -2, -3])
+        raises(ValueError, "bincount(b)")
+        c = array([0, 1, 2])
+        w = array([1, 2])
+        raises(ValueError, "bincount(c, w)")
+        raises(ValueError, "bincount([])")
+