[pypy-commit] pypy numpy-ufunc-object: give ufuncs a reduce method, make stuff use it.

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: numpy-ufunc-object
Changeset: r46925:b78512479c78
Date: 2011-08-30 14:36 -0400
http://bitbucket.org/pypy/pypy/changeset/b78512479c78/

Log:	give ufuncs a reduce method, make stuff use it.

diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -20,6 +20,7 @@
 
 class FakeSpace(object):
     w_ValueError = None
+    w_TypeError = None
 
     def __init__(self):
         """NOT_RPYTHON"""
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
@@ -90,30 +90,13 @@
     descr_rpow = _binop_right_impl("power")
     descr_rmod = _binop_right_impl("mod")
 
-    def _reduce_sum_prod_impl(op_name, init):
-        reduce_driver = jit.JitDriver(greens=['signature'],
-                         reds = ['i', 'size', 'self', 'result', 'res_dtype'])
+    def _reduce_ufunc_impl(ufunc_name):
+        def impl(self, space):
+            return getattr(interp_ufuncs.get(space), ufunc_name).descr_reduce(space, self)
+        return func_with_new_name(impl, "reduce_%s_impl" % ufunc_name)
 
-        def loop(self, res_dtype, result, size):
-            i = 0
-            while i < size:
-                reduce_driver.jit_merge_point(signature=self.signature,
-                                              self=self, res_dtype=res_dtype,
-                                              size=size, i=i, result=result)
-                result = getattr(res_dtype, op_name)(
-                    result,
-                    self.eval(i).convert_to(res_dtype)
-                )
-                i += 1
-            return result
-
-        def impl(self, space):
-            dtype = interp_ufuncs.find_unaryop_result_dtype(
-                space, self.find_dtype(), promote_to_largest=True
-            )
-            result = dtype.adapt_val(init)
-            return loop(self, dtype, result, self.find_size()).wrap(space)
-        return func_with_new_name(impl, "reduce_%s_impl" % op_name)
+    descr_sum = _reduce_ufunc_impl("add")
+    descr_prod = _reduce_ufunc_impl("multiply")
 
     def _reduce_max_min_impl(op_name):
         reduce_driver = jit.JitDriver(greens=['signature'],
@@ -192,8 +175,6 @@
     def descr_any(self, space):
         return space.wrap(self._any())
 
-    descr_sum = _reduce_sum_prod_impl("add", 0)
-    descr_prod = _reduce_sum_prod_impl("mul", 1)
     descr_max = _reduce_max_min_impl("max")
     descr_min = _reduce_max_min_impl("min")
     descr_argmax = _reduce_argmax_argmin_impl("max")
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -3,9 +3,15 @@
 from pypy.interpreter.gateway import interp2app
 from pypy.interpreter.typedef import TypeDef, GetSetProperty, interp_attrproperty
 from pypy.module.micronumpy import interp_dtype, signature
+from pypy.rlib import jit
 from pypy.tool.sourcetools import func_with_new_name
 
 
+reduce_driver = jit.JitDriver(
+    greens = ["signature"],
+    reds = ["i", "size", "self", "dtype", "value", "obj"]
+)
+
 class W_Ufunc(Wrappable):
     def __init__(self, name, promote_to_float, promote_bools, identity):
         self.name = name
@@ -29,6 +35,39 @@
             raise OperationError(space.w_TypeError, space.wrap(str(e)))
         return self.call(space, args_w)
 
+    def descr_reduce(self, space, w_obj):
+        from pypy.module.micronumpy.interp_numarray import convert_to_array, Scalar
+
+        if self.argcount != 2:
+            raise OperationError(space.w_ValueError, space.wrap("reduce only supported for binary functions"))
+        if self.identity is None:
+            raise OperationError(space.w_NotImplementedError, space.wrap("%s is missing its identity value" % self.name))
+
+        obj = convert_to_array(space, w_obj)
+        if isinstance(obj, Scalar):
+            raise OperationError(space.w_TypeError, space.wrap("cannot reduce on a scalar"))
+
+        size = obj.find_size()
+        dtype = find_unaryop_result_dtype(
+            space, obj.find_dtype(),
+            promote_to_largest=True
+        )
+        value = self.identity.convert_to(dtype)
+        new_sig = signature.Signature.find_sig([
+            self.reduce_signature, obj.signature
+        ])
+        return self.reduce(new_sig, value, obj, dtype, size).wrap(space)
+
+    def reduce(self, signature, value, obj, dtype, size):
+        i = 0
+        while i < size:
+            reduce_driver.jit_merge_point(signature=signature, self=self,
+                                          value=value, obj=obj, i=i,
+                                          dtype=dtype, size=size)
+            value = self.func(dtype, value, obj.eval(i).convert_to(dtype))
+            i += 1
+        return value
+
 class W_Ufunc1(W_Ufunc):
     argcount = 1
 
@@ -68,6 +107,7 @@
         W_Ufunc.__init__(self, name, promote_to_float, promote_bools, identity)
         self.func = func
         self.signature = signature.Call2(func)
+        self.reduce_signature = signature.BaseSignature()
 
     def call(self, space, args_w):
         from pypy.module.micronumpy.interp_numarray import (Call2,
@@ -100,7 +140,9 @@
     __repr__ = interp2app(W_Ufunc.descr_repr),
 
     identity = GetSetProperty(W_Ufunc.descr_get_identity),
-    nin = interp_attrproperty("argcount", cls=W_Ufunc)
+    nin = interp_attrproperty("argcount", cls=W_Ufunc),
+
+    reduce = interp2app(W_Ufunc.descr_reduce),
 )
 
 def find_binop_result_dtype(space, dt1, dt2, promote_to_float=False,
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -298,3 +298,13 @@
         b = arctan(a)
         assert math.isnan(b[0])
 
+    def test_reduce_errors(self):
+        from numpy import sin, add
+
+        raises(ValueError, sin.reduce, [1, 2, 3])
+        raises(TypeError, add.reduce, 1)
+
+    def test_reduce(self):
+        from numpy import add
+
+        assert add.reduce([1, 2, 3]) == 6
\ No newline at end of file