[pypy-svn] r46423 - in pypy/dist/pypy/rpython/numpy: . test

simonb at codespeak.net simonb at codespeak.net
Sat Sep 8 23:20:37 CEST 2007


Author: simonb
Date: Sat Sep  8 23:20:36 2007
New Revision: 46423

Modified:
   pypy/dist/pypy/rpython/numpy/rarray.py
   pypy/dist/pypy/rpython/numpy/test/test_array.py
Log:
basic numpy arithmetic working.

Modified: pypy/dist/pypy/rpython/numpy/rarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/rarray.py	(original)
+++ pypy/dist/pypy/rpython/numpy/rarray.py	Sat Sep  8 23:20:36 2007
@@ -195,15 +195,6 @@
         it.dataptr[0] = value
         iter_next(it)
 
-def ll_array_add(it0, it1, it2, iter_next):
-    debug_assert(it0.size == it1.size, "it0.size == it1.size")
-    debug_assert(it1.size == it2.size, "it0.size == it1.size")
-    while it0.index < it0.size:
-        it0.dataptr[0] = it1.dataptr[0] + it2.dataptr[0]
-        iter_next(it0)
-        iter_next(it1)
-        iter_next(it2)
-
 def dim_of_ITER(ITER):
     return ITER.TO.coordinates.length
 
@@ -310,41 +301,57 @@
         key = self.__class__, self.typecode, self.ndim
         return key
 
+def ll_array_binop(it0, it1, it2, iter_next, binop):
+    debug_assert(it0.size == it1.size, "it0.size == it1.size")
+    debug_assert(it1.size == it2.size, "it0.size == it1.size")
+    while it0.index < it0.size:
+        it0.dataptr[0] = binop(it1.dataptr[0], it2.dataptr[0])
+        iter_next(it0)
+        iter_next(it1)
+        iter_next(it2)
 
-class __extend__(pairtype(ArrayRepr, ArrayRepr)):
-    def rtype_add((r_array1, r_array2), hop):
-        v_array1, v_array2 = hop.inputargs(r_array1, r_array2)
-        r_array0 = hop.r_result
-        cARRAY = hop.inputconst(Void, r_array0.ARRAY.TO)
-        # We build a contiguous "result" array
-        # from the largest of the two args:
-        v_bigarray = hop.gendirectcall(ll_find_largest, cARRAY, v_array1, v_array2)
-        v_array0 = hop.gendirectcall(ll_build_like, cARRAY, v_bigarray)
-        iter_new, iter_reset, iter_broadcast, iter_next = gen_iter_funcs(r_array0.ndim)
-        cITER = hop.inputconst(Void, r_array0.ITER.TO)
-        creset = hop.inputconst(Void, iter_reset)
-        cbroadcast = hop.inputconst(Void, iter_broadcast)
-        cnext = hop.inputconst(Void, iter_next)
-        v_it0 = hop.gendirectcall(iter_new, cITER, v_array0, v_bigarray, creset, cbroadcast)
-        v_it1 = hop.gendirectcall(iter_new, cITER, v_array1, v_bigarray, creset, cbroadcast)
-        v_it2 = hop.gendirectcall(iter_new, cITER, v_array2, v_bigarray, creset, cbroadcast)
-        return hop.gendirectcall(ll_array_add, v_it0, v_it1, v_it2, cnext)
-
-class __extend__(pairtype(ArrayRepr, Repr)):
-    def rtype_add((r_array, r_ob), hop):
-        assert 0
-        v_ob = hop.inputarg(r_ob, 1)
-        if isinstance(r_ob.lowleveltype, Primitive):
-            r_item, v_item = convert_scalar_to_array(r_array, v_ob, hop.llops)
-
-class __extend__(pairtype(Repr, ArrayRepr)):
-    def rtype_add((r_ob, r_array), hop):
-        # XX ach! how to get this to work ??
-        assert 0
-        return pair(r_array, r_ob).rtype_add(hop)
-
+def _rtype_binop(r_array0, r_array1, r_array2, v_array1, v_array2, hop, binop):
+    cARRAY = hop.inputconst(Void, r_array0.ARRAY.TO)
+    # We build a contiguous "result" array
+    # from the largest of the two args:
+    v_array0 = hop.gendirectcall(ll_build_like, cARRAY, v_array1, v_array2)
+    iter_new, iter_reset, iter_broadcast, iter_next = gen_iter_funcs(r_array0.ndim)
+    cITER = hop.inputconst(Void, r_array0.ITER.TO)
+    creset = hop.inputconst(Void, iter_reset)
+    cbroadcast = hop.inputconst(Void, iter_broadcast)
+    cnext = hop.inputconst(Void, iter_next)
+    v_it0 = hop.gendirectcall(iter_new, cITER, v_array0, v_array0, creset, cbroadcast)
+    v_it1 = hop.gendirectcall(iter_new, cITER, v_array1, v_array0, creset, cbroadcast)
+    v_it2 = hop.gendirectcall(iter_new, cITER, v_array2, v_array0, creset, cbroadcast)
+    cbinop = hop.inputconst(Void, binop)
+    hop.gendirectcall(ll_array_binop, v_it0, v_it1, v_it2, cnext, cbinop)
+    return v_array0
+
+def rtype_binop((r_array1, r_array2), hop, binop):
+    r_array0 = hop.r_result
+    v_array1, v_array2 = hop.inputargs(r_array1, r_array2)
+
+    if isinstance(r_array1.lowleveltype, Primitive):
+        r_array1, v_array1 = convert_scalar_to_array(r_array0, v_array1, hop.llops)
+    elif not isinstance(r_array1, ArrayRepr):
+        raise TyperError("can't operate with %s"%r_array1)
+
+    if isinstance(r_array2.lowleveltype, Primitive):
+        r_array2, v_array2 = convert_scalar_to_array(r_array0, v_array2, hop.llops)
+    elif not isinstance(r_array2, ArrayRepr):
+        raise TyperError("can't operate with %s"%r_array2)
+    print "rtype_binop", binop, binop(2,3)
+    return _rtype_binop(r_array0, r_array1, r_array2, v_array1, v_array2, hop, binop)
+
+for tp in (pairtype(ArrayRepr, ArrayRepr),
+           pairtype(ArrayRepr, Repr),
+           pairtype(Repr, ArrayRepr)):
+    for (binop, methname) in ((lambda a,b:a+b, "rtype_add"),
+                              (lambda a,b:a-b, "rtype_sub"),
+                              (lambda a,b:a*b, "rtype_mul"),
+                              (lambda a,b:a/b, "rtype_div")):
+        setattr(tp, methname, lambda self, hop, binop=binop : rtype_binop(self, hop, binop))
 
-        
 def gen_getset_item(ndim):
     unrolling_dims = unrolling_iterable(range(ndim))
     def ll_get_item(ARRAY, ao, tpl):
@@ -447,7 +454,6 @@
     s_array = r_array.s_array.get_one_dim()
     r_array = llops.rtyper.getrepr(s_array)
     from pypy.rpython.rmodel import inputconst
-#    ARRAY, _ = ArrayRepr.make_types(1, ITEM)
     cARRAY = inputconst(Void, r_array.ARRAY.TO)
     cITEM = inputconst(Void, r_array.ITEM)
     v_casted = llops.genop("cast_primitive", [v_item], r_array.ITEM)
@@ -602,47 +608,47 @@
     array.dataptr = ao.dataptr
     return array
 
-def ll_find_largest(ARRAY, array0, array1):
+def ll_build_like(ARRAY, array0, array1):
+    # Build with shape from the largest of array0 or array1.
+    # Note we cannot take the union of array0 and array1.
+    ndim = max(array0.ndim, array1.ndim)
+    array = ll_allocate(ARRAY, ndim)
     sz0 = ll_mul_list(array0.shape, array0.ndim)
     sz1 = ll_mul_list(array1.shape, array1.ndim)
-    # XXX 
-    if sz0 > sz1:
-        return array0
-    return array1
-
-def ll_build_like(ARRAY, ao):
-    array = ll_allocate(ARRAY, ao.ndim)
-    sz = ll_mul_list(ao.shape)
+    sz = max(sz0, sz1)
     array.data = malloc(ARRAY.data.TO, sz)
-    array.dataptr = array.data
+    array.dataptr = direct_arrayitems(array.data)
     itemsize = 1
-    i = ao.ndim - 1
+    i = ndim - 1
     while i >= 0:
-        size = ao.shape[i]
+        if sz0>sz1:
+            size = array0.shape[i]
+        else:
+            size = array1.shape[i]
         array.shape[i] = size
         array.strides[i] = itemsize
         itemsize *= size
         i -= 1
     return array
 
-def ll_setitem1(array, index, item):
-    array.data[index] = item
-
-def ll_getitem1(array, index):
-    return array.data[index]
-
-def ll_add(ARRAY, a1, a2):
-    size = a1.shape[0]
-    if size != a2.shape[0]:
-        raise ValueError
-    array = malloc(ARRAY)
-    array.data = malloc(ARRAY.data.TO, size)
-    i = 0
-    while i < size:
-        array.data[i] = a1.data[i] + a2.data[i]
-        i += 1
-    array.dataptr = direct_arrayitems(array.data)
-    return array
+#def ll_setitem1(array, index, item):
+#    array.data[index] = item
+#
+#def ll_getitem1(array, index):
+#    return array.data[index]
+
+#def ll_add(ARRAY, a1, a2):
+#    size = a1.shape[0]
+#    if size != a2.shape[0]:
+#        raise ValueError
+#    array = malloc(ARRAY)
+#    array.data = malloc(ARRAY.data.TO, size)
+#    i = 0
+#    while i < size:
+#        array.data[i] = a1.data[i] + a2.data[i]
+#        i += 1
+#    array.dataptr = direct_arrayitems(array.data)
+#    return array
 
 def ll_transpose(ARRAY, ao):
     ndim = ao.ndim

Modified: pypy/dist/pypy/rpython/numpy/test/test_array.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/test/test_array.py	(original)
+++ pypy/dist/pypy/rpython/numpy/test/test_array.py	Sat Sep  8 23:20:36 2007
@@ -106,7 +106,6 @@
         assert s.ndim == 2
 
     def test_annotate_array_add(self):
-        py.test.skip("broken ATM")
         def f():
             a1 = numpy.array([1,2])
             a2 = numpy.array([6,9])
@@ -118,7 +117,6 @@
         assert s.typecode == 'i'
 
     def test_annotate_array_add_coerce(self):
-        py.test.skip("broken ATM")
         def f():
             a1 = numpy.array([1,2])
             a2 = numpy.array([6.,9.])
@@ -130,7 +128,6 @@
         assert s.typecode == 'd'
 
     def test_annotate_array_add_2(self):
-        py.test.skip("broken ATM")
         def f():
             a = numpy.array([1,2])
             a = a + 3
@@ -280,25 +277,6 @@
         res = interpret(access_with_variable, [])
         assert res == 45
 
-    def test_specialize_array_add(self):
-        py.test.skip("broken ATM")
-        def f():
-            a1 = numpy.array([1.,2.])
-            a2 = numpy.array([6,9])
-            return a1 + a2
-
-        res = interpret(f, [])
-        assert res.data[0] == 7
-        assert res.data[1] == 11
-
-        def f():
-            a1 = numpy.array([1,2])
-            return a1 + 2
-
-        res = interpret(f, [])
-        assert res.data[0] == 3
-        assert res.data[1] == 4
-
     def test_specialize_array_attr(self):
         def f():
             a = numpy.array([1,2])
@@ -388,6 +366,29 @@
         for i in range(len(data)):
             assert res.dataptr[i] == data[i]
 
+    def test_specialize_list_coerce(self):
+        py.test.skip('coercion not implemented')
+        def f():
+            a = numpy.zeros((3,4), dtype='i')
+            a[:] = [3.,4.,4.,2.]
+            return a
+        res = interpret(f, [])
+        data = [3,4,4,2]*3
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
+    def test_specialize_rhs_coerce(self):
+        py.test.skip('coercion not implemented')
+        def f():
+            a = numpy.zeros((4,), dtype='i')
+            b = numpy.array([3.,4.,4.,2.])
+            a[:] = b
+            return a
+        res = interpret(f, [])
+        data = [3,4,4,2]
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
     def test_specialize_slice_1_0(self):
         def f():
             a = numpy.zeros((12,), dtype='i')
@@ -518,7 +519,101 @@
         from pypy.translator.backendopt.test.test_malloc import TestLLTypeMallocRemoval
         TestLLTypeMallocRemoval.check_malloc_removed(graph)
 
+    def test_specialize_array_add_1_0(self):
+        def f():
+            a1 = numpy.array(range(4,10))
+            a2 = numpy.array([3])
+            return a1 + a2
+        data = [i+3 for i in range(4,10)]
+        res = interpret(f, [])
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
+        def f():
+            a = numpy.array(range(4,10))
+            return a + 3
+        res = interpret(f, [])
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
 
+        def f():
+            a = numpy.array(range(4,10))
+            return 3 + a
+        res = interpret(f, [])
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
+    def test_specialize_array_sub_1_0(self):
+        def f():
+            a1 = numpy.array(range(4,10))
+            a2 = numpy.array([3])
+            return a1 - a2
+        data = [i-3 for i in range(4,10)]
+        res = interpret(f, [])
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
+        def f():
+            a = numpy.array(range(4,10))
+            return a - 3
+        res = interpret(f, [])
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
+        def f():
+            a = numpy.array(range(4,10))
+            return 3 - a
+        data = [3-i for i in range(4,10)]
+        res = interpret(f, [])
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+
+
+    def test_specialize_array_add_1_1(self):
+        def f():
+            a1 = numpy.array([1,2])
+            a2 = numpy.array([6,9])
+            return a1 + a2
+
+        res = interpret(f, [])
+        assert res.data[0] == 7
+        assert res.data[1] == 11
+
+    def test_specialize_array_add_1_1_coerce(self):
+        py.test.skip('coercion not implemented')
+        def f():
+            a1 = numpy.array([1,2])
+            a2 = numpy.array([6.,9.])
+            return a1 + a2
+
+        res = interpret(f, [])
+        assert res.data[0] == 7.
+        assert res.data[1] == 11.
+
+    def test_specialize_array_add_2_1(self):
+        def f():
+            a1 = numpy.array([1,2,3,4]).reshape((2,2))
+            a2 = numpy.array([5,6])
+            return a1 + a2
+
+        res = interpret(f, [])
+        data = [6,8,8,10]
+        for i in range(len(data)):
+            assert res.dataptr[i] == data[i]
+        assert res.shape[0] == 2
+        assert res.shape[1] == 2
+        assert res.strides[0] == 2
+        assert res.strides[1] == 1
+
+    def test_specialize_array_mul_0_2(self):
+        def f():
+            a = numpy.array([1,2,3,4]).reshape((2,2))
+            return 5*a
+
+        res = interpret(f, [])
+        data = [1,2,3,4]
+        for i in range(len(data)):
+            assert res.dataptr[i] == 5*data[i]
 
 class Test_compile:
     def setup_class(self):



More information about the Pypy-commit mailing list