[pypy-commit] pypy numpy-impicit-convert: Convert sources of ufuncs to numarrays if needed

[snus_mumrik]

Author: Ilya Osadchiy <osadchiy.ilya at 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