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