Author: Maciej Fijalkowski <[email protected]>
Branch: 
Changeset: r45478:516f5a4af65d
Date: 2011-07-12 00:31 +0200
http://bitbucket.org/pypy/pypy/changeset/516f5a4af65d/

Log:    (justinpeel, fijal) Reduce-like functions for micronumpy

diff --git a/pypy/module/micronumpy/interp_numarray.py 
b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -1,5 +1,5 @@
 from pypy.interpreter.baseobjspace import ObjSpace, W_Root, Wrappable
-from pypy.interpreter.error import operationerrfmt
+from pypy.interpreter.error import operationerrfmt, OperationError
 from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.interpreter.typedef import TypeDef, GetSetProperty
 from pypy.rlib import jit
@@ -49,6 +49,14 @@
     return math.pow(v1, v2)
 def mod(v1, v2):
     return math.fmod(v1, v2)
+def maximum(v1, v2):
+    return max(v1, v2)
+def minimum(v1, v2):
+    return min(v1, v2)
+def and_bool(v1, v2):
+    return v1 and bool(v2)
+def or_bool(v1, v2):
+    return v1 or bool(v2)
 
 class BaseArray(Wrappable):
     def __init__(self):
@@ -110,6 +118,129 @@
     descr_pow = _binop_impl(pow)
     descr_mod = _binop_impl(mod)
 
+    def _reduce_sum_prod_impl(function, init):
+        reduce_driver = jit.JitDriver(greens=['signature'],
+                         reds = ['i', 'size', 'self', 'result'])
+
+        def loop(self, result, size):
+            i = 0
+            while i < size:
+                reduce_driver.jit_merge_point(signature=self.signature,
+                                              self=self, size=size, i=i,
+                                              result=result)
+                result = function(result, self.eval(i))
+                i += 1
+            return result
+
+        def impl(self, space):
+            return space.wrap(loop(self, init, self.find_size()))
+        return func_with_new_name(impl, "reduce_%s_impl" % function.__name__)
+
+    def _reduce_max_min_impl(function):
+        reduce_driver = jit.JitDriver(greens=['signature'],
+                         reds = ['i', 'size', 'self', 'result'])    
+        def loop(self, result, size):
+            i = 1
+            while i < size:
+                reduce_driver.jit_merge_point(signature=self.signature,
+                                              self=self, size=size, i=i,
+                                              result=result)
+                result = function(result, self.eval(i))
+                i += 1
+            return result
+        
+        def impl(self, space):
+            size = self.find_size()
+            if size == 0:
+                raise OperationError(space.w_ValueError,
+                    space.wrap("Can't call %s on zero-size arrays" \
+                            % function.__name__))
+            return space.wrap(loop(self, self.eval(0), size))
+        return func_with_new_name(impl, "reduce_%s_impl" % function.__name__)
+
+    def _reduce_argmax_argmin_impl(function):
+        reduce_driver = jit.JitDriver(greens=['signature'],
+                         reds = ['i', 'size', 'result', 'self'])
+        def loop(self, size):
+            result = 0
+            cur_best = self.eval(0)
+            i = 1
+            while i < size:
+                reduce_driver.jit_merge_point(signature=self.signature,
+                                              self=self, size=size, i=i,
+                                              result=result)
+                new_best = function(cur_best, self.eval(i))
+                if new_best != cur_best:
+                    result = i
+                    cur_best = new_best
+                i += 1
+            return result
+        def impl(self, space):
+            size = self.find_size()
+            if size == 0:
+                raise OperationError(space.w_ValueError,
+                    space.wrap("Can't call %s on zero-size arrays" \
+                            % function.__name__))
+            return space.wrap(loop(self, size))
+        return func_with_new_name(impl, "reduce_arg%s_impl" % 
function.__name__)
+
+    def _reduce_all_impl():
+        reduce_driver = jit.JitDriver(greens=['signature'],
+                         reds = ['i', 'size', 'result', 'self'])
+        def loop(self, result, size):
+            i = 0
+            while i < size:
+                reduce_driver.jit_merge_point(signature=self.signature,
+                                              self=self, size=size, i=i,
+                                              result=result)
+                result = and_bool(result, self.eval(i))
+                if not result:
+                    break
+                i += 1
+            return result
+
+        def impl(self, space):
+            size = self.find_size()
+            return space.wrap(loop(self, True, size))
+        return func_with_new_name(impl, "reduce_all_impl")
+
+    def _reduce_any_impl():
+        reduce_driver = jit.JitDriver(greens=['signature'],
+                         reds = ['i', 'size', 'result', 'self'])       
+
+        def loop(self, result, size):
+            i = 0
+            while i < size:
+                reduce_driver.jit_merge_point(signature=self.signature,
+                                              self=self, size=size, i=i,
+                                              result=result)
+                result = or_bool(result, self.eval(i))
+                if result:
+                    break
+                i += 1
+            return result
+        def impl(self, space):
+            size = self.find_size()
+            return space.wrap(loop(self, False, size))
+        return func_with_new_name(impl, "reduce_any_impl")
+
+    descr_sum = _reduce_sum_prod_impl(add, 0.0)
+    descr_prod = _reduce_sum_prod_impl(mul, 1.0)
+    descr_max = _reduce_max_min_impl(maximum)
+    descr_min = _reduce_max_min_impl(minimum)
+    descr_argmax = _reduce_argmax_argmin_impl(maximum)
+    descr_argmin = _reduce_argmax_argmin_impl(minimum)
+    descr_all = _reduce_all_impl()
+    descr_any = _reduce_any_impl()
+
+    def descr_dot(self, space, w_other):
+        if isinstance(w_other, BaseArray):
+            w_res = self.descr_mul(space, w_other)
+            assert isinstance(w_res, BaseArray)
+            return w_res.descr_sum(space)
+        else:
+            return self.descr_mul(space, w_other)
+
     def get_concrete(self):
         raise NotImplementedError
 
@@ -136,13 +267,7 @@
         return self.get_concrete().descr_setitem(space, item, value)
 
     def descr_mean(self, space):
-        s = 0
-        concrete = self.get_concrete()
-        size = concrete.find_size()
-        for i in xrange(size):
-            s += concrete.getitem(i)
-        return space.wrap(s / size)
-
+        return 
space.wrap(space.float_w(self.descr_sum(space))/self.find_size())
 
 class FloatWrapper(BaseArray):
     """
@@ -391,4 +516,13 @@
     __mod__ = interp2app(BaseArray.descr_mod),
 
     mean = interp2app(BaseArray.descr_mean),
+    sum = interp2app(BaseArray.descr_sum),
+    prod = interp2app(BaseArray.descr_prod),
+    max = interp2app(BaseArray.descr_max),
+    min = interp2app(BaseArray.descr_min),
+    argmax = interp2app(BaseArray.descr_argmax),
+    argmin = interp2app(BaseArray.descr_argmin),
+    all = interp2app(BaseArray.descr_all),
+    any = interp2app(BaseArray.descr_any),
+    dot = interp2app(BaseArray.descr_dot),
 )
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
@@ -282,12 +282,88 @@
         assert a.mean() == 2.0
         assert a[:4].mean() == 1.5
 
+    def test_sum(self):
+        from numpy import array
+        a = array(range(5))
+        assert a.sum() == 10.0
+        assert a[:4].sum() == 6.0
+
+    def test_prod(self):
+        from numpy import array
+        a = array(range(1,6))
+        assert a.prod() == 120.0
+        assert a[:4].prod() == 24.0
+
+    def test_max(self):
+        from numpy import array
+        a = array([-1.2, 3.4, 5.7, -3.0, 2.7])
+        assert a.max() == 5.7
+        b = array([])
+        raises(ValueError, "b.max()")
+
+    def test_max_add(self):
+        from numpy import array
+        a = array([-1.2, 3.4, 5.7, -3.0, 2.7])
+        assert (a+a).max() == 11.4
+
+    def test_min(self):
+        from numpy import array
+        a = array([-1.2, 3.4, 5.7, -3.0, 2.7])
+        assert a.min() == -3.0
+        b = array([])
+        raises(ValueError, "b.min()")
+
+    def test_argmax(self):
+        from numpy import array
+        a = array([-1.2, 3.4, 5.7, -3.0, 2.7])
+        assert a.argmax() == 2
+        b = array([])
+        raises(ValueError, "b.argmax()")
+
+    def test_argmin(self):
+        from numpy import array
+        a = array([-1.2, 3.4, 5.7, -3.0, 2.7])
+        assert a.argmin() == 3
+        b = array([])
+        raises(ValueError, "b.argmin()")
+
+    def test_all(self):
+        from numpy import array
+        a = array(range(5))
+        assert a.all() == False
+        a[0] = 3.0
+        assert a.all() == True
+        b = array([])
+        assert b.all() == True
+
+    def test_any(self):
+        from numpy import array, zeros
+        a = array(range(5))
+        assert a.any() == True
+        b = zeros(5)
+        assert b.any() == False
+        c = array([])
+        assert c.any() == False
+
+    def test_dot(self):
+        from numpy import array
+        a = array(range(5))
+        assert a.dot(a) == 30.0
+
+    def test_dot_constant(self):
+        from numpy import array
+        a = array(range(5))
+        b = a.dot(2.5)
+        for i in xrange(5):
+            assert b[i] == 2.5*a[i]
+
+
 class AppTestSupport(object):
     def setup_class(cls):
         import struct
         cls.space = gettestobjspace(usemodules=('micronumpy',))
         cls.w_data = cls.space.wrap(struct.pack('dddd', 1, 2, 3, 4))
-    
+
     def test_fromstring(self):
         from numpy import fromstring
         a = fromstring(self.data)
diff --git a/pypy/module/micronumpy/test/test_zjit.py 
b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -4,9 +4,13 @@
     FloatWrapper, Call2, SingleDimSlice, add, mul, neg, Call1)
 from pypy.module.micronumpy.interp_ufuncs import negative
 from pypy.module.micronumpy.compile import numpy_compile
+from pypy.rlib.objectmodel import specialize
 
 class FakeSpace(object):
-    pass
+    w_ValueError = None
+    @specialize.argtype(1)
+    def wrap(self, v):
+        return v
 
 class TestNumpyJIt(LLJitMixin):
     def setup_class(cls):
@@ -51,6 +55,93 @@
 
         assert result == f(5)
 
+    def test_sum(self):
+        space = self.space
+
+        def f(i):
+            ar = SingleDimArray(i)
+            return ar.descr_add(space, ar).descr_sum(space)
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({"getarrayitem_raw": 2, "float_add": 2,
+                          "int_add": 1,
+                          "int_lt": 1, "guard_true": 1, "jump": 1})
+        assert result == f(5)
+
+    def test_prod(self):
+        space = self.space
+
+        def f(i):
+            ar = SingleDimArray(i)
+            return ar.descr_add(space, ar).descr_prod(space)
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({"getarrayitem_raw": 2, "float_add": 1,
+                          "float_mul": 1, "int_add": 1,
+                          "int_lt": 1, "guard_true": 1, "jump": 1})
+        assert result == f(5)
+
+    def test_max(self):
+        space = self.space
+
+        def f(i):
+            ar = SingleDimArray(i)
+            j = 0
+            while j < i:
+                ar.get_concrete().storage[j] = float(j)
+                j += 1
+            return ar.descr_add(space, ar).descr_max(space)
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({"getarrayitem_raw": 2, "float_add": 1,
+                          "float_gt": 1, "int_add": 1,
+                          "int_lt": 1, "guard_true": 1, 
+                          "guard_false": 1, "jump": 1})
+
+    def test_min(self):
+        space = self.space
+
+        def f(i):
+            ar = SingleDimArray(i)
+            j = 0
+            while j < i:
+                ar.get_concrete().storage[j] = float(j)
+                j += 1
+            return ar.descr_add(space, ar).descr_min(space)
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({"getarrayitem_raw": 2, "float_add": 1,
+                           "float_lt": 1, "int_add": 1,
+                           "int_lt": 1, "guard_true": 2,
+                           "jump": 1})
+
+    def test_all(self):
+        space = self.space
+
+        def f(i):
+            ar = SingleDimArray(i)
+            j = 0
+            while j < i:
+                ar.get_concrete().storage[j] = 1.0
+                j += 1
+            return ar.descr_add(space, ar).descr_all(space)
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({"getarrayitem_raw": 2, "float_add": 1,
+                          "int_add": 1, "float_ne": 1,
+                          "int_lt": 1, "guard_true": 2, "jump": 1})
+
+    def test_any(self):
+        space = self.space
+
+        def f(i):
+            ar = SingleDimArray(i)
+            return ar.descr_add(space, ar).descr_any(space)
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({"getarrayitem_raw": 2, "float_add": 1,
+                          "int_add": 1, "float_ne": 1, "guard_false": 1,
+                          "int_lt": 1, "guard_true": 1, "jump": 1})
+
     def test_already_forecd(self):
         def f(i):
             ar = SingleDimArray(i)
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to