Author: Alex Gaynor <alex.gay...@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 _______________________________________________ pypy-commit mailing list pypy-commit@python.org http://mail.python.org/mailman/listinfo/pypy-commit