Author: Maciej Fijalkowski <fij...@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()
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
http://mail.python.org/mailman/listinfo/pypy-commit
