Author: Ilya Osadchiy <osadchiy.i...@gmail.com> Branch: numpy-impicit-convert Changeset: r45091:abdff8d681cc Date: 2011-06-23 22:14 +0300 http://bitbucket.org/pypy/pypy/changeset/abdff8d681cc/
Log: Convert sources of ufuncs to numarrays if needed 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 OperationError, operationerrfmt from pypy.interpreter.gateway import interp2app, unwrap_spec from pypy.interpreter.typedef import TypeDef, GetSetProperty from pypy.rlib import jit @@ -114,6 +114,14 @@ s += concrete.getitem(i) return space.wrap(s / size) +def access_as_array (space, w_obj): + try: + # If it's a scalar + return FloatWrapper(space.float_w(w_obj)) + except OperationError: + # Convert to array. + # Could we somehow use COW in some cases? + return new_numarray(space, w_obj) class FloatWrapper(BaseArray): """ @@ -321,14 +329,17 @@ def __del__(self): lltype.free(self.storage, flavor='raw') -def descr_new_numarray(space, w_type, w_size_or_iterable): +def new_numarray(space, w_size_or_iterable): l = space.listview(w_size_or_iterable) arr = SingleDimArray(len(l)) i = 0 for w_elem in l: arr.storage[i] = space.float_w(space.float(w_elem)) i += 1 - return space.wrap(arr) + return arr + +def descr_new_numarray(space, w_type, w_size_or_iterable): + return space.wrap(new_numarray(space, w_size_or_iterable)) @unwrap_spec(size=int) def zeros(space, size): 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 @@ -1,10 +1,15 @@ import math from pypy.interpreter.gateway import unwrap_spec -from pypy.module.micronumpy.interp_numarray import BaseArray, Call1, Call2, Signature +from pypy.module.micronumpy.interp_numarray import BaseArray, Call1, Call2, Signature, access_as_array from pypy.rlib import rfloat from pypy.tool.sourcetools import func_with_new_name +def _issequence(space, w_obj): + # Copied from cpyext's PySequence_Check + """Return True if the object provides sequence protocol, and False otherwise. + This function always succeeds.""" + return (space.findattr(w_obj, space.wrap("__getitem__")) is not None) def ufunc(func): signature = Signature() @@ -13,19 +18,45 @@ w_res = Call1(func, w_obj, w_obj.signature.transition(signature)) w_obj.invalidates.append(w_res) return w_res - return space.wrap(func(space.float_w(w_obj))) + elif _issequence(space, w_obj): + w_obj_arr = access_as_array(space, w_obj) + w_res = Call1(func, w_obj_arr, w_obj_arr.signature.transition(signature)) + return w_res + else: + return space.wrap(func(space.float_w(w_obj))) return func_with_new_name(impl, "%s_dispatcher" % func.__name__) def ufunc2(func): signature = Signature() def impl(space, w_lhs, w_rhs): - if isinstance(w_lhs, BaseArray) and isinstance(w_rhs, BaseArray): + lhs_is_array = isinstance(w_lhs, BaseArray) + rhs_is_array = isinstance(w_rhs, BaseArray) + if lhs_is_array and rhs_is_array: + # This is the (most likely) fall-through case in conversion checks + # Not sure if making it a special case makes it much faster new_sig = w_lhs.signature.transition(signature).transition(w_rhs.signature) w_res = Call2(func, w_lhs, w_rhs, new_sig) w_lhs.invalidates.append(w_res) w_rhs.invalidates.append(w_res) return w_res - return space.wrap(func(space.float_w(w_lhs), space.float_w(w_rhs))) + elif _issequence(space, w_lhs) or _issequence(space, w_rhs): + if lhs_is_array: + w_lhs_arr = w_lhs + else: + w_lhs_arr = access_as_array(space, w_lhs) + if rhs_is_array: + w_rhs_arr = w_rhs + else: + w_rhs_arr = access_as_array(space, w_rhs) + new_sig = w_lhs_arr.signature.transition(signature).transition(w_rhs_arr.signature) + w_res = Call2(func, w_lhs_arr, w_rhs_arr, new_sig) + if lhs_is_array: + w_lhs_arr.invalidates.append(w_res) + if rhs_is_array: + w_rhs_arr.invalidates.append(w_res) + return w_res + else: + return space.wrap(func(space.float_w(w_lhs), space.float_w(w_rhs))) return func_with_new_name(impl, "%s_dispatcher" % func.__name__) @ufunc 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 @@ -10,6 +10,28 @@ assert sign(-0.0) == 0.0 assert minimum(2.0, 3.0) == 2.0 + def test_sequence(self): + from numpy import array, negative, minimum + a = array(range(3)) + b = [2.0, 1.0, 0.0] + c = 1.0 + b_neg = negative(b) + assert isinstance(b_neg, array) + for i in range(3): + assert b_neg[i] == -b[i] + min_a_b = minimum(a, b) + assert isinstance(min_a_b, array) + for i in range(3): + assert min_a_b[i] == min(a[i], b[i]) + min_a_c = minimum(a, c) + assert isinstance(min_a_c, array) + for i in range(3): + assert min_a_c[i] == min(a[i], c) + min_b_c = minimum(b, c) + assert isinstance(min_b_c, array) + for i in range(3): + assert min_b_c[i] == min(b[i], c) + def test_negative(self): from numpy import array, negative _______________________________________________ pypy-commit mailing list pypy-commit@python.org http://mail.python.org/mailman/listinfo/pypy-commit