Author: Maciej Fijalkowski <fij...@gmail.com> Branch: Changeset: r45793:fb245ad56f79 Date: 2011-07-21 09:54 +0200 http://bitbucket.org/pypy/pypy/changeset/fb245ad56f79/
Log: merge numpy-slice 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 @@ -23,6 +23,8 @@ reds = ['result_size', 'i', 'self', 'result']) all_driver = jit.JitDriver(greens=['signature'], reds=['i', 'size', 'self']) any_driver = jit.JitDriver(greens=['signature'], reds=['i', 'size', 'self']) +slice_driver1 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'self', 'arr']) +slice_driver2 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'self', 'arr']) class Signature(object): def __init__(self): @@ -288,10 +290,18 @@ res = SingleDimSlice(start, stop, step, slice_length, self, self.signature.transition(SingleDimSlice.static_signature)) return space.wrap(res) - @unwrap_spec(item=int, value=float) - def descr_setitem(self, space, item, value): + def descr_setitem(self, space, w_idx, w_value): + # TODO: indexing by tuples and lists self.invalidated() - return self.get_concrete().descr_setitem(space, item, value) + start, stop, step, slice_length = space.decode_index4(w_idx, + self.find_size()) + if step == 0: + # Single index + self.get_concrete().setitem(start, + space.float_w(w_value)) + else: + self.get_concrete().setslice(space, start, stop, step, + slice_length, w_value) def descr_mean(self, space): return space.wrap(space.float_w(self.descr_sum(space))/self.find_size()) @@ -440,8 +450,8 @@ return self.parent.getitem(self.calc_index(item)) @unwrap_spec(item=int, value=float) - def descr_setitem(self, space, item, value): - return self.parent.descr_setitem(space, self.calc_index(item), value) + def setitem(self, item, value): + return self.parent.setitem(self.calc_index(item), value) def descr_len(self, space): return space.wrap(self.find_size()) @@ -455,14 +465,58 @@ def __init__(self, start, stop, step, slice_length, parent, signature): ViewArray.__init__(self, parent, signature) + if isinstance(parent, SingleDimSlice): + self.start = parent.calc_index(start) + self.stop = parent.calc_index(stop) + self.step = parent.step * step + self.parent = parent.parent + else: self.start = start self.stop = stop self.step = step + self.parent = parent self.size = slice_length def find_size(self): return self.size + def _sliceloop1(self, start, stop, step, arr): + storage = self.parent.storage + signature = Signature() + new_sig = self.signature.transition(signature) + i = start + j = 0 + while i < stop: + slice_driver1.jit_merge_point(signature=signature, self=self, + step=step, stop=stop, i=i, j=j, arr=arr) + storage[i] = arr.eval(j) + j += 1 + i += step + + def _sliceloop2(self, start, stop, step, arr): + storage = self.parent.storage + signature = Signature() + new_sig = self.signature.transition(signature) + i = start + j = 0 + while i > stop: + slice_driver2.jit_merge_point(signature=signature, self=self, + step=step, stop=stop, i=i, j=j, arr=arr) + storage[i] = arr.eval(j) + j += 1 + i += step + + def setslice(self, space, start, stop, step, slice_length, arr): + arr = convert_to_array(space, arr) + start = self.calc_index(start) + if stop != -1: + stop = self.calc_index(stop) + step = self.step * step + if step > 0: + self._sliceloop1(start, stop, step, arr) + else: + self._sliceloop2(start, stop, step, arr) + def calc_index(self, item): return (self.start + item * self.step) @@ -486,7 +540,7 @@ def eval(self, i): return self.storage[i] - def getindex(self, space, item): + def getindex(self, item): if item >= self.size: raise operationerrfmt(space.w_IndexError, '%d above array size', item) @@ -504,11 +558,44 @@ return self.storage[item] @unwrap_spec(item=int, value=float) - def descr_setitem(self, space, item, value): - item = self.getindex(space, item) + def setitem(self, item, value): + item = self.getindex(item) self.invalidated() self.storage[item] = value + def _sliceloop1(self, start, stop, step, arr): + signature = Signature() + new_sig = self.signature.transition(signature) + i = start + j = 0 + while i < stop: + slice_driver1.jit_merge_point(signature=signature, self=self, + step=step, stop=stop, i=i, j=j, arr=arr) + self.storage[i] = arr.eval(j) + j += 1 + i += step + + def _sliceloop2(self, start, stop, step, arr): + signature = Signature() + new_sig = self.signature.transition(signature) + i = start + j = 0 + while i > stop: + slice_driver2.jit_merge_point(signature=signature, self=self, + step=step, stop=stop, i=i, j=j, arr=arr) + self.storage[i] = arr.eval(j) + j += 1 + i += step + + def setslice(self, space, start, stop, step, slice_length, arr): + i = start + if not isinstance(arr, BaseArray): + arr = convert_to_array(space, arr) + if step > 0: + self._sliceloop1(start, stop, step, arr) + else: + self._sliceloop2(start, stop, step, arr) + def __del__(self): lltype.free(self.storage, flavor='raw') 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 @@ -92,6 +92,48 @@ raises(IndexError, "a[5] = 0.0") raises(IndexError, "a[-6] = 3.0") + def test_setslice_array(self): + from numpy import array + a = array(range(5)) + b = array(range(2)) + a[1:4:2] = b + assert a[1] == 0. + assert a[3] == 1. + + def test_setslice_of_slice_array(self): + from numpy import array, zeros + a = zeros(5) + a[::2] = array([9., 10., 11.]) + assert a[0] == 9. + assert a[2] == 10. + assert a[4] == 11. + a[1:4:2][::-1] = array([1., 2.]) + assert a[0] == 9. + assert a[1] == 2. + assert a[2] == 10. + assert a[3] == 1. + assert a[4] == 11. + a = zeros(10) + a[::2][::-1][::2] = array(range(1,4)) + a[8] = 1. + a[4] = 2. + a[0] = 3. + + def test_setslice_list(self): + from numpy import array + a = array(range(5)) + b = [0., 1.] + a[1:4:2] = b + assert a[1] == 0. + assert a[3] == 1. + + def test_setslice_constant(self): + from numpy import array + a = array(range(5)) + a[1:4:2] = 0. + assert a[1] == 0. + assert a[3] == 0. + def test_len(self): from numpy import array a = array(range(5)) diff --git a/pypy/module/micronumpy/test/test_zjit.py b/pypy/module/micronumpy/test/test_zjit.py --- a/pypy/module/micronumpy/test/test_zjit.py +++ b/pypy/module/micronumpy/test/test_zjit.py @@ -5,6 +5,7 @@ from pypy.module.micronumpy.interp_ufuncs import negative from pypy.module.micronumpy.compile import numpy_compile from pypy.rlib.objectmodel import specialize +from pypy.rlib.nonconst import NonConstant class FakeSpace(object): w_ValueError = None @@ -248,6 +249,24 @@ 'int_lt': 1, 'guard_true': 1, 'jump': 1}) assert result == f(5) + def test_setslice(self): + space = self.space + + def f(i): + step = NonConstant(3) + ar = SingleDimArray(step*i) + ar2 = SingleDimArray(i) + ar2.storage[1] = 5.5 + ar.setslice(space, 0, step*i, step, i, ar2.descr_add(space, ar2)) + return ar.get_concrete().storage[3] + + result = self.meta_interp(f, [5], listops=True, backendopt=True) + self.check_loops({'getarrayitem_raw': 2, + 'float_add' : 1, + 'setarrayitem_raw': 1, 'int_add': 2, + 'int_lt': 1, 'guard_true': 1, 'jump': 1}) + assert result == 11.0 + class TestTranslation(object): def test_compile(self): x = numpy_compile('aa+f*f/a-', 10) _______________________________________________ pypy-commit mailing list pypy-commit@python.org http://mail.python.org/mailman/listinfo/pypy-commit