[pypy-commit] pypy numpypy-axisops: rework AxisIterator, make single tests pass

mattip noreply at buildbot.pypy.org
Tue Dec 27 21:48:35 CET 2011


Author: mattip
Branch: numpypy-axisops
Changeset: r50914:b12a872f0961
Date: 2011-12-27 07:32 +0200
http://bitbucket.org/pypy/pypy/changeset/b12a872f0961/

Log:	rework AxisIterator, make single tests pass

diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -111,17 +111,19 @@
     """ This object will return offsets of each start of a stride on the 
         desired dimension, starting at the desired index
     """
-    def __init__(self, start, strides, backstrides, shape, dim=-1, start=[]):
+    def __init__(self, arr_start, strides, backstrides, shape, dim=-1, slice_start=[]):
         self.shape = shape
-        self.indices = [0] * len(arr.shape)
+        self.indices = [0] * len(shape)
         self.done = False
-        self.offset = start
+        self.offset = arr_start
         self.dim = len(shape) - 1
+        self.strides = strides
+        self.backstrides = backstrides
         if dim >= 0:
             self.dim = dim
-        if len(start) == len(shape):
-            for i in range(len(start)):
-                self.offset += strides[i] * start[i]
+        if len(slice_start) == len(shape):
+            for i in range(len(slice_start)):
+                self.offset += strides[i] * slice_start[i]
     def next(self, shapelen):
         offset = self.offset
         indices = [0] * shapelen
@@ -146,6 +148,6 @@
         res.backstrides = self.backstrides
         res.shape = self.shape
         res.dim = self.dim
-        res.done = done
+        res.done = self.done
         return res
 
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
@@ -9,7 +9,7 @@
 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib.rstring import StringBuilder
 from pypy.module.micronumpy.interp_iter import ArrayIterator,\
-     view_iter_from_arr, OneDimIterator, AxisIterator
+     view_iter_from_arr, OneDimIterator, axis_iter_from_arr
 
 numpy_driver = jit.JitDriver(
     greens=['shapelen', 'sig'],
@@ -280,6 +280,8 @@
     
     def _reduce_ufunc_impl(ufunc_name):
         def impl(self, space, w_dim=None):
+            if w_dim is None:
+                w_dim = space.wrap(w_dim)
             return getattr(interp_ufuncs.get(space), ufunc_name).reduce(space,
                                                        self, True, w_dim)
         return func_with_new_name(impl, "reduce_%s_impl" % ufunc_name)
@@ -756,7 +758,7 @@
         shapelen = len(result.shape)
         sig = self.find_sig()
         ri = ArrayIterator(self.size)
-        si = AxisIterator(self,self.dim)
+        si = axis_iter_from_arr(self, self.dim)
         while not ri.done():
             frame = sig.create_frame(self, self.values, chunks = si.indices)
             val = sig.eval(frame, self)
@@ -985,23 +987,24 @@
     def _fast_setslice(self, space, w_value):
         assert isinstance(w_value, ConcreteArray)
         itemsize = self.dtype.itemtype.get_element_size()
-        if len(self.shape) == 1:
+        shapelen = len(self.shape)
+        if shapelen == 1:
             rffi.c_memcpy(
                 rffi.ptradd(self.storage, self.start * itemsize),
                 rffi.ptradd(w_value.storage, w_value.start * itemsize),
                 self.size * itemsize
             )
         else:
-            dest = AxisIterator(self)
-            source = AxisIterator(w_value)
+            dest = axis_iter_from_arr(self)
+            source = axis_iter_from_arr(w_value)
             while not dest.done:
                 rffi.c_memcpy(
                     rffi.ptradd(self.storage, dest.offset * itemsize),
                     rffi.ptradd(w_value.storage, source.offset * itemsize),
                     self.shape[-1] * itemsize
                 )
-                source.next()
-                dest.next()
+                source = source.next(shapelen)
+                dest = dest.next(shapelen)
 
     def _sliceloop(self, source, res_shape):
         sig = source.find_sig(res_shape)
diff --git a/pypy/module/micronumpy/test/test_iterators.py b/pypy/module/micronumpy/test/test_iterators.py
--- a/pypy/module/micronumpy/test/test_iterators.py
+++ b/pypy/module/micronumpy/test/test_iterators.py
@@ -8,24 +8,31 @@
 
 class TestAxisIteratorDirect(object):
     def test_axis_iterator(self):
+        a = W_NDimArray(5*3, [5, 3], MockDtype(), 'C')
+        i = AxisIterator(a)
+        ret = []
+        while not i.done:
+            ret.append(i.offset)
+            i = i.next()
+        assert ret == [0, 3, 6, 9, 12]
         a = W_NDimArray(7*5*3, [7, 5, 3], MockDtype(), 'C')
         i = AxisIterator(a)
         ret = []
         while not i.done:
             ret.append(i.offset)
-            i.next()
+            i = i.next()
         assert ret == [3*v for v in range(7*5)]
         i = AxisIterator(a,2)
         ret = []
         while not i.done:
             ret.append(i.offset)
-            i.next()
+            i = i.next()
         assert ret == [3*v for v in range(7*5)]
         i = AxisIterator(a,1)
         ret = []
         while not i.done:
             ret.append(i.offset)
-            i.next()
+            i = i.next()
         assert ret == [ 0,  1,  2, 15, 16, 17, 30, 31, 32, 45, 46, 47,
                        60, 61, 62, 75, 76, 77, 90, 91, 92] 
     def test_axis_iterator_with_start(self):
@@ -34,18 +41,18 @@
         ret = []
         while not i.done:
             ret.append(i.offset)
-            i.next()
+            i = i.next()
         assert ret == [3*v for v in range(7*5)]
         i = AxisIterator(a, start=[1, 1, 0])
         ret = []
         while not i.done:
             ret.append(i.offset)
-            i.next()
+            i = i.next()
         assert ret == [3*v+18 for v in range(7*5)]
         i = AxisIterator(a, 1, [2, 0, 2])
         ret = []
         while not i.done:
             ret.append(i.offset)
-            i.next()
+            i = i.next()
         assert ret == [v + 32 for v in [ 0,  1,  2, 15, 16, 17, 30, 31, 32,
                             45, 46, 47, 60, 61, 62, 75, 76, 77, 90, 91, 92]]
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -720,6 +720,9 @@
         assert a.sum() == 5
 
         raises(TypeError, 'a.sum(2, 3)')
+
+    def test_sumND(self):
+        skip('Not finished yet')
         a = arange(15).reshape(5, 3)
         assert (a.sum(0) == [30, 35, 40]).all()
         assert (a.sum(1) == [3, 12, 21, 30, 39]).all()


More information about the pypy-commit mailing list