[pypy-commit] pypy numpy-refactor: next_skip_x and port test_iter

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-refactor
Changeset: r57184:afefaf666d03
Date: 2012-09-06 18:42 +0200
http://bitbucket.org/pypy/pypy/changeset/afefaf666d03/

Log:	next_skip_x and port test_iter

diff --git a/pypy/module/micronumpy/arrayimpl/concrete.py b/pypy/module/micronumpy/arrayimpl/concrete.py
--- a/pypy/module/micronumpy/arrayimpl/concrete.py
+++ b/pypy/module/micronumpy/arrayimpl/concrete.py
@@ -29,6 +29,10 @@
     def next(self):
         self.offset += self.skip
 
+    def next_skip_x(self, x):
+        self.offset += self.skip * x
+        self.index += x
+
     def done(self):
         return self.offset >= self.size
 
@@ -45,6 +49,10 @@
         self.offset += self.skip
         self.index += 1
 
+    def next_skip_x(self, x):
+        self.offset += self.skip * x
+        self.index += x
+
     def done(self):
         return self.index >= self.size
 
@@ -74,6 +82,22 @@
             self._done = True
         self.offset = offset
 
+    @jit.unroll_safe
+    def next_skip_x(self, step):
+        for i in range(len(self.shape) - 1, -1, -1):
+            if self.indexes[i] < self.shape[i] - step:
+                self.indexes[i] += step
+                self.offset += self.strides[i] * step
+                break
+            else:
+                remaining_step = (self.indexes[i] + step) // self.shape[i]
+                this_i_step = step - remaining_step * self.shape[i]
+                self.offset += self.strides[i] * this_i_step
+                self.indexes[i] = self.indexes[i] +  this_i_step
+                step = remaining_step
+        else:
+            self._done = True
+
     def done(self):
         return self._done
 
diff --git a/pypy/module/micronumpy/interp_flatiter.py b/pypy/module/micronumpy/interp_flatiter.py
--- a/pypy/module/micronumpy/interp_flatiter.py
+++ b/pypy/module/micronumpy/interp_flatiter.py
@@ -6,8 +6,8 @@
 
 class W_FlatIterator(Wrappable):
     def __init__(self, arr):
-        self.arr = arr
-        self.iter = self.arr.create_iter()
+        self.base = arr
+        self.iter = arr.create_iter()
         self.index = 0
 
     def descr_next(self, space):
@@ -18,6 +18,38 @@
         self.index += 1
         return w_res
 
+    @jit.unroll_safe
+    def descr_getitem(self, space, w_idx):
+        if not (space.isinstance_w(w_idx, space.w_int) or
+            space.isinstance_w(w_idx, space.w_slice)):
+            raise OperationError(space.w_IndexError,
+                                 space.wrap('unsupported iterator index'))
+        base = self.base
+        start, stop, step, length = space.decode_index4(w_idx, base.get_size())
+        # setslice would have been better, but flat[u:v] for arbitrary
+        # shapes of array a cannot be represented as a[x1:x2, y1:y2]
+        base_iter = base.create_iter()
+        xxx
+        return base.getitem(basei.offset)
+        base_iter = ViewIterator(base.start, base.strides,
+                             base.backstrides, base.shape)
+        shapelen = len(base.shape)
+        basei = basei.next_skip_x(shapelen, start)
+        res = W_NDimArray([lngth], base.dtype, base.order)
+        ri = res.create_iter()
+        while not ri.done():
+            flat_get_driver.jit_merge_point(shapelen=shapelen,
+                                             base=base,
+                                             basei=basei,
+                                             step=step,
+                                             res=res,
+                                             ri=ri)
+            w_val = base.getitem(basei.offset)
+            res.setitem(ri.offset, w_val)
+            basei = basei.next_skip_x(shapelen, step)
+            ri = ri.next(shapelen)
+        return res
+
     def descr_iter(self):
         return self
 
diff --git a/pypy/module/micronumpy/test/test_iter.py b/pypy/module/micronumpy/test/test_iter.py
--- a/pypy/module/micronumpy/test/test_iter.py
+++ b/pypy/module/micronumpy/test/test_iter.py
@@ -1,4 +1,4 @@
-from pypy.module.micronumpy.interp_iter import ViewIterator
+from pypy.module.micronumpy.arrayimpl.concrete import MultiDimViewIterator
 
 class TestIterDirect(object):
     def test_C_viewiterator(self):
@@ -9,36 +9,36 @@
         strides = [5, 1]
         backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
         assert backstrides == [10, 4]
-        i = ViewIterator(start, strides, backstrides, shape)
-        i = i.next(2)
-        i = i.next(2)
-        i = i.next(2)
+        i = MultiDimViewIterator(None, start, strides, backstrides, shape)
+        i.next()
+        i.next()
+        i.next()
         assert i.offset == 3
         assert not i.done()
-        assert i.indices == [0,3]
+        assert i.indexes == [0,3]
         #cause a dimension overflow
-        i = i.next(2)
-        i = i.next(2)
+        i.next()
+        i.next()
         assert i.offset == 5
-        assert i.indices == [1,0]
+        assert i.indexes == [1,0]
 
         #Now what happens if the array is transposed? strides[-1] != 1
         # therefore layout is non-contiguous
         strides = [1, 3]
         backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
         assert backstrides == [2, 12]
-        i = ViewIterator(start, strides, backstrides, shape)
-        i = i.next(2)
-        i = i.next(2)
-        i = i.next(2)
+        i = MultiDimViewIterator(None, start, strides, backstrides, shape)
+        i.next()
+        i.next()
+        i.next()
         assert i.offset == 9
         assert not i.done()
-        assert i.indices == [0,3]
+        assert i.indexes == [0,3]
         #cause a dimension overflow
-        i = i.next(2)
-        i = i.next(2)
+        i.next()
+        i.next()
         assert i.offset == 1
-        assert i.indices == [1,0]
+        assert i.indexes == [1,0]
 
     def test_C_viewiterator_step(self):
         #iteration in C order with #contiguous layout => strides[-1] is 1
@@ -48,22 +48,22 @@
         strides = [5, 1]
         backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
         assert backstrides == [10, 4]
-        i = ViewIterator(start, strides, backstrides, shape)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
+        i = MultiDimViewIterator(None, start, strides, backstrides, shape)
+        i.next_skip_x(2)
+        i.next_skip_x(2)
+        i.next_skip_x(2)
         assert i.offset == 6
         assert not i.done()
-        assert i.indices == [1,1]
+        assert i.indexes == [1,1]
         #And for some big skips
-        i = i.next_skip_x(2,5)
+        i.next_skip_x(5)
         assert i.offset == 11
-        assert i.indices == [2,1]
-        i = i.next_skip_x(2,5)
+        assert i.indexes == [2,1]
+        i.next_skip_x(5)
         # Note: the offset does not overflow but recycles,
         # this is good for broadcast
         assert i.offset == 1
-        assert i.indices == [0,1]
+        assert i.indexes == [0,1]
         assert i.done()
 
         #Now what happens if the array is transposed? strides[-1] != 1
@@ -71,18 +71,18 @@
         strides = [1, 3]
         backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
         assert backstrides == [2, 12]
-        i = ViewIterator(start, strides, backstrides, shape)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
+        i = MultiDimViewIterator(None, start, strides, backstrides, shape)
+        i.next_skip_x(2)
+        i.next_skip_x(2)
+        i.next_skip_x(2)
         assert i.offset == 4
-        assert i.indices == [1,1]
+        assert i.indexes == [1,1]
         assert not i.done()
-        i = i.next_skip_x(2,5)
+        i.next_skip_x(5)
         assert i.offset == 5
-        assert i.indices == [2,1]
+        assert i.indexes == [2,1]
         assert not i.done()
-        i = i.next_skip_x(2,5)
-        assert i.indices == [0,1]
+        i.next_skip_x(5)
+        assert i.indexes == [0,1]
         assert i.offset == 3
         assert i.done()