Author: Maciej Fijalkowski <[email protected]>
Branch: numpy-indexing-by-arrays-2
Changeset: r51516:d1cbcc34848d
Date: 2012-01-20 12:20 +0200
http://bitbucket.org/pypy/pypy/changeset/d1cbcc34848d/
Log: first approximation of array[array-of-int]
diff --git a/pypy/module/micronumpy/interp_boxes.py
b/pypy/module/micronumpy/interp_boxes.py
--- a/pypy/module/micronumpy/interp_boxes.py
+++ b/pypy/module/micronumpy/interp_boxes.py
@@ -104,7 +104,10 @@
_attrs_ = ()
class W_IntegerBox(W_NumberBox):
- pass
+ def convert_to_int(self):
+ from pypy.rpython.lltypesystem import rffi, lltype
+
+ return rffi.cast(lltype.Signed, self.value)
class W_SignedIntegerBox(W_IntegerBox):
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
@@ -20,6 +20,9 @@
if self.step != 0:
shape.append(self.lgt)
+ def get_iter(self):
+ xxx
+
class IntArrayChunk(BaseChunk):
def __init__(self, arr):
self.arr = arr.get_concrete()
@@ -27,10 +30,22 @@
def extend_shape(self, shape):
shape.extend(self.arr.shape)
+ def get_iter(self):
+ return self.arr.create_iter()
+
+ def get_index(self, iter):
+ return self.arr.getitem(iter.offset).convert_to_int()
+
class BoolArrayChunk(BaseChunk):
def __init__(self, arr):
self.arr = arr.get_concrete()
+ def extend_shape(self, shape):
+ xxx
+
+ def get_iter(self):
+ xxx
+
class BaseTransform(object):
pass
@@ -212,6 +227,30 @@
return self._done
# ------ other iterators that are not part of the computation frame ----------
+
+class ChunkIterator(object):
+ def __init__(self, shape, chunks):
+ self.chunks = chunks
+ self.indices = [0] * len(shape)
+ self.shape = shape
+ self.chunk_iters = [chunk.get_iter() for chunk in self.chunks]
+
+ def next(self, shapelen):
+ for i in range(shapelen - 1, -1, -1):
+ if self.indices[i] < self.shape[i] - 1:
+ self.indices[i] += 1
+ self.chunk_iters[i] = self.chunk_iters[i].next()
+ break
+ else:
+ self.indices[i] = 0
+ # XXX reset one dim iter probably
+ return self
+
+ def get_index(self, shapelen):
+ l = []
+ for i in range(shapelen):
+ l.append(self.chunks[i].get_index(self.chunk_iters[i]))
+ return l
class SkipLastAxisIterator(object):
def __init__(self, arr):
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
@@ -10,7 +10,8 @@
from pypy.tool.sourcetools import func_with_new_name
from pypy.rlib.rstring import StringBuilder
from pypy.module.micronumpy.interp_iter import ArrayIterator, OneDimIterator,\
- SkipLastAxisIterator, Chunk, ViewIterator, BoolArrayChunk, IntArrayChunk
+ SkipLastAxisIterator, Chunk, ViewIterator, BoolArrayChunk, IntArrayChunk,\
+ ChunkIterator
numpy_driver = jit.JitDriver(
greens=['shapelen', 'sig'],
@@ -583,6 +584,9 @@
item = concrete._index_of_single_item(space, w_idx)
return concrete.getitem(item)
chunks = self._prepare_slice_args(space, w_idx)
+ for chunk in chunks:
+ if not isinstance(chunk, Chunk):
+ return self.get_concrete().force_slice_getitem(space, chunks)
return space.wrap(self.create_slice(chunks))
def descr_setitem(self, space, w_idx, w_value):
@@ -604,13 +608,6 @@
view = self.create_slice(chunks).get_concrete()
view.setslice(space, w_value)
- def force_slice(self, shape, chunks):
- size = 1
- for elem in shape:
- size *= elem
- res = W_NDimArray(size, shape, self.find_dtype())
- xxx
-
@jit.unroll_safe
def create_slice(self, chunks):
shape = []
@@ -620,9 +617,6 @@
s = i + 1
assert s >= 0
shape += self.shape[s:]
- for chunk in chunks:
- if not isinstance(chunk, Chunk):
- return self.force_slice(shape, chunks)
if not isinstance(self, ConcreteArray):
return VirtualSlice(self, chunks, shape)
r = calculate_slice_strides(self.shape, self.start, self.strides,
@@ -1104,6 +1098,20 @@
builder.append(']')
@jit.unroll_safe
+ def _index_of_single_item_int(self, space, index):
+ item = self.start
+ for i in range(len(index)):
+ v = index[i]
+ if v < 0:
+ v += self.shape[i]
+ if v < 0 or v >= self.shape[i]:
+ raise operationerrfmt(space.w_IndexError,
+ "index (%d) out of range (0<=index<%d", i, self.shape[i],
+ )
+ item += v * self.strides[i]
+ return item
+
+ @jit.unroll_safe
def _index_of_single_item(self, space, w_idx):
if space.isinstance_w(w_idx, space.w_int):
idx = space.int_w(w_idx)
@@ -1115,17 +1123,7 @@
return self.start + idx * self.strides[0]
index = [space.int_w(w_item)
for w_item in space.fixedview(w_idx)]
- item = self.start
- for i in range(len(index)):
- v = index[i]
- if v < 0:
- v += self.shape[i]
- if v < 0 or v >= self.shape[i]:
- raise operationerrfmt(space.w_IndexError,
- "index (%d) out of range (0<=index<%d", i, self.shape[i],
- )
- item += v * self.strides[i]
- return item
+ return self._index_of_single_item_int(space, index)
def setslice(self, space, w_value):
res_shape = shape_agreement(space, self.shape, w_value.shape)
@@ -1178,6 +1176,28 @@
def fill(self, space, w_value):
self.setslice(space, scalar_w(space, self.dtype, w_value))
+ def force_slice_getitem(self, space, chunks):
+ shape = []
+ i = -1
+ for i, chunk in enumerate(chunks):
+ chunk.extend_shape(shape)
+ s = i + 1
+ assert s >= 0
+ shape += self.shape[s:]
+ size = 1
+ for elem in shape:
+ size *= elem
+ res = W_NDimArray(size, shape, self.find_dtype())
+ ri = res.create_iter()
+ ci = ChunkIterator(shape, chunks)
+ shapelen = len(shape)
+ while not ri.done():
+ index = ci.get_index(shapelen)
+ v = self.getitem(self._index_of_single_item_int(space, index))
+ res.setitem(ri.offset, v)
+ ri = ri.next(shapelen)
+ ci = ci.next(shapelen)
+ return res
class ViewArray(ConcreteArray):
def create_sig(self):
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
