[pypy-commit] pypy default: merge numpy-slice

fijal noreply at buildbot.pypy.org
Thu Jul 21 09:55:41 CEST 2011


Author: Maciej Fijalkowski <fijall at 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)


More information about the pypy-commit mailing list