[pypy-commit] pypy numpy-setslice: numpy: setslice added but doesn't work correctly for setting a slice of a slice yet

justinpeel noreply at buildbot.pypy.org
Sat Jul 16 20:49:15 CEST 2011


Author: Justin Peel <notmuchtotell at gmail.com>
Branch: numpy-setslice
Changeset: r45685:9d55cebe8768
Date: 2011-07-16 12:49 -0600
http://bitbucket.org/pypy/pypy/changeset/9d55cebe8768/

Log:	numpy: setslice added but doesn't work correctly for setting a slice
	of a slice yet

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
@@ -21,6 +21,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):
@@ -255,10 +257,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())
@@ -407,8 +417,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())
@@ -430,6 +440,47 @@
     def find_size(self):
         return self.size
 
+    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.parent.setitem(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.parent.setitem(i, arr.eval(j))
+            j += 1
+            i += step
+
+    def setslice(self, space, start, stop, step, slice_length, arr):
+        # can't set a slice of a slice yet
+        if stop < 0:
+            stop += self.find_size()
+        if step > 0:
+            stop = min(stop, self.find_size())
+        else:
+            stop = max(stop, 0)
+        arr = convert_to_array(space, arr)
+        start = self.calc_index(start)
+        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)
 
@@ -453,7 +504,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)
@@ -471,11 +522,50 @@
         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 stop < 0:
+            stop += self.find_size()
+        if step > 0:
+            stop = min(stop, self.find_size())
+        else:
+            stop = max(stop, 0)
+        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
@@ -60,6 +60,34 @@
         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.
+        # a[1:4:2][::-1] = b # does not work yet
+        c=a[1:4:2][::-1]
+        c[:] = b
+        assert a[1] == 1.
+        assert a[3] == 0.
+
+    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,21 @@
                           '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)
+            ar.setslice(space, 0, step*i, step, i, ar2)
+            return ar.get_concrete().storage[3]
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({'getarrayitem_raw': 1,
+                          'setarrayitem_raw': 1, 'int_add': 2,
+                          'int_lt': 1, 'guard_true': 1, 'jump': 1})
+
 class TestTranslation(object):
     def test_compile(self):
         x = numpy_compile('aa+f*f/a-', 10)


More information about the pypy-commit mailing list