Author: Justin Peel <[email protected]>
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)
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
