Author: mattip <matti.pi...@gmail.com>
Branch: 
Changeset: r82336:d77888929462
Date: 2016-02-19 16:29 +0200
http://bitbucket.org/pypy/pypy/changeset/d77888929462/

Log:    test, fix indexing ndarray with scalar, single boolean array

diff --git a/pypy/module/micronumpy/concrete.py 
b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -12,8 +12,8 @@
     ArrayArgumentException, W_NumpyObject
 from pypy.module.micronumpy.iterators import ArrayIter
 from pypy.module.micronumpy.strides import (
-    IntegerChunk, SliceChunk, NewAxisChunk, EllipsisChunk, new_view,
-    calc_strides, calc_new_strides, shape_agreement,
+    IntegerChunk, SliceChunk, NewAxisChunk, EllipsisChunk, BooleanChunk,
+    new_view, calc_strides, calc_new_strides, shape_agreement,
     calculate_broadcast_strides, calc_backstrides, calc_start, is_c_contiguous,
     is_f_contiguous)
 from rpython.rlib.objectmodel import keepalive_until_here
@@ -236,6 +236,8 @@
 
     @jit.unroll_safe
     def _prepare_slice_args(self, space, w_idx):
+        print '_prepare_slice_args', w_idx
+        from pypy.module.micronumpy import boxes
         if space.isinstance_w(w_idx, space.w_str):
             raise oefmt(space.w_IndexError, "only integers, slices (`:`), "
                 "ellipsis (`...`), numpy.newaxis (`None`) and integer or "
@@ -258,6 +260,7 @@
         result = []
         i = 0
         has_ellipsis = False
+        has_filter = False
         for w_item in space.fixedview(w_idx):
             if space.is_w(w_item, space.w_Ellipsis):
                 if has_ellipsis:
@@ -272,6 +275,16 @@
             elif space.isinstance_w(w_item, space.w_slice):
                 result.append(SliceChunk(w_item))
                 i += 1
+            elif isinstance(w_item, W_NDimArray) and 
w_item.get_dtype().is_bool():
+                if has_filter:
+                    # in CNumPy, the support for this is incomplete
+                    raise oefmt(space.w_ValueError,
+                        "an index can only have a single boolean mask; "
+                        "use np.take or create a sinlge mask array")
+                has_filter = True
+                result.append(BooleanChunk(w_item))
+            elif isinstance(w_item, boxes.W_GenericBox):
+                result.append(IntegerChunk(w_item.descr_int(space)))
             else:
                 result.append(IntegerChunk(w_item))
                 i += 1
@@ -280,11 +293,14 @@
         return result
 
     def descr_getitem(self, space, orig_arr, w_index):
+        print 'concrete descr_gettiem %s' % str(w_index)[:35]
         try:
             item = self._single_item_index(space, w_index)
+            print 'concrete descr_gettiem _single_item_index succeeded'
             return self.getitem(item)
         except IndexError:
             # not a single result
+            print 'concrete descr_gettiem _single_item_index failed'
             chunks = self._prepare_slice_args(space, w_index)
             return new_view(space, orig_arr, chunks)
 
diff --git a/pypy/module/micronumpy/ndarray.py 
b/pypy/module/micronumpy/ndarray.py
--- a/pypy/module/micronumpy/ndarray.py
+++ b/pypy/module/micronumpy/ndarray.py
@@ -107,8 +107,9 @@
             arr = W_NDimArray(self.implementation.transpose(self, None))
         return space.wrap(loop.tostring(space, arr))
 
-    def getitem_filter(self, space, arr):
-        if arr.ndims() > 1 and arr.get_shape() != self.get_shape():
+    def getitem_filter(self, space, arr, axis=0):
+        shape = self.get_shape()
+        if arr.ndims() > 1 and arr.get_shape() != shape:
             raise OperationError(space.w_IndexError, space.wrap(
                 "boolean index array should have 1 dimension"))
         if arr.get_size() > self.get_size():
@@ -116,14 +117,14 @@
                 "index out of range for array"))
         size = loop.count_all_true(arr)
         if arr.ndims() == 1:
-            if self.ndims() > 1 and arr.get_shape()[0] != self.get_shape()[0]:
+            if self.ndims() > 1 and arr.get_shape()[0] != shape[axis]:
                 msg = ("boolean index did not match indexed array along"
-                      " dimension 0; dimension is %d but corresponding"
-                      " boolean dimension is %d" % (self.get_shape()[0],
+                      " dimension %d; dimension is %d but corresponding"
+                      " boolean dimension is %d" % (axis, shape[axis],
                       arr.get_shape()[0]))
                 #warning = 
space.gettypefor(support.W_VisibleDeprecationWarning)
                 space.warn(space.wrap(msg), space.w_VisibleDeprecationWarning)
-            res_shape = [size] + self.get_shape()[1:]
+            res_shape = shape[:axis] + [size] + shape[axis+1:]
         else:
             res_shape = [size]
         w_res = W_NDimArray.from_shape(space, res_shape, self.get_dtype(),
@@ -149,6 +150,8 @@
     def _prepare_array_index(self, space, w_index):
         if isinstance(w_index, W_NDimArray):
             return [], w_index.get_shape(), w_index.get_shape(), [w_index]
+        if isinstance(w_index, boxes.W_GenericBox):
+            return [], [1], [1], [w_index]
         w_lst = space.listview(w_index)
         for w_item in w_lst:
             if not (space.isinstance_w(w_item, space.w_int) or 
space.isinstance_w(w_item, space.w_float)):
@@ -162,7 +165,14 @@
         arr_index_in_shape = False
         prefix = []
         for i, w_item in enumerate(w_lst):
-            if (isinstance(w_item, W_NDimArray) or
+            if isinstance(w_item, W_NDimArray) and 
w_item.get_dtype().is_bool():
+                if w_item.ndims() > 0:
+                    indexes_w[i] = w_item
+                else:
+                    raise oefmt(space.w_IndexError,
+                        "in the future, 0-d boolean arrays will be "
+                        "interpreted as a valid boolean index")
+            elif (isinstance(w_item, W_NDimArray) or
                     space.isinstance_w(w_item, space.w_list)):
                 w_item = convert_to_array(space, w_item)
                 if shape is None:
@@ -232,6 +242,8 @@
                 raise oefmt(space.w_IndexError,
                         "in the future, 0-d boolean arrays will be "
                         "interpreted as a valid boolean index")
+        elif isinstance(w_idx, boxes.W_GenericBox):
+            w_ret = self.getitem_array_int(space, w_idx)
         else:
             try:
                 w_ret = self.implementation.descr_getitem(space, self, w_idx)
diff --git a/pypy/module/micronumpy/strides.py 
b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -77,14 +77,40 @@
         backstride = base_stride * max(0, base_length - 1)
         return 0, base_length, base_stride, backstride
 
+class BooleanChunk(BaseChunk):
+    input_dim = 1
+    out_dim = 1
+    def __init__(self, w_idx):
+        self.w_idx = w_idx
+
+    def compute(self, space, base_length, base_stride):
+        raise oefmt(space.w_NotImplementedError, 'cannot reach')
 
 def new_view(space, w_arr, chunks):
     arr = w_arr.implementation
-    r = calculate_slice_strides(space, arr.shape, arr.start, arr.get_strides(),
-                                arr.get_backstrides(), chunks)
+    dim = -1
+    for i, c in enumerate(chunks):
+        if isinstance(c, BooleanChunk):
+            dim = i
+            break
+    if dim >= 0:
+        # filter by axis r
+        filtr = chunks.pop(dim)
+        assert isinstance(filtr, BooleanChunk) 
+        w_arr = w_arr.getitem_filter(space, filtr.w_idx, axis=dim)
+        arr = w_arr.implementation
+        r = calculate_slice_strides(space, arr.shape, arr.start,
+                 arr.get_strides(), arr.get_backstrides(), chunks)
+    else:
+        r = calculate_slice_strides(space, arr.shape, arr.start,
+                     arr.get_strides(), arr.get_backstrides(), chunks)
     shape, start, strides, backstrides = r
-    return W_NDimArray.new_slice(space, start, strides[:], backstrides[:],
+    w_ret = W_NDimArray.new_slice(space, start, strides[:], backstrides[:],
                                  shape[:], arr, w_arr)
+    if dim == 0:
+        # Do not return a view
+        return w_ret.descr_copy(space, space.wrap(w_ret.get_order()))
+    return w_ret
 
 @jit.unroll_safe
 def _extend_shape(old_shape, chunks):
@@ -127,7 +153,7 @@
                      jit.isconstant(len(chunks)))
 def calculate_slice_strides(space, shape, start, strides, backstrides, chunks):
     """
-    Note: `chunks` must contain exactly one EllipsisChunk object.
+    Note: `chunks` can contain at most one EllipsisChunk object.
     """
     size = 0
     used_dims = 0
diff --git a/pypy/module/micronumpy/test/test_deprecations.py 
b/pypy/module/micronumpy/test/test_deprecations.py
--- a/pypy/module/micronumpy/test/test_deprecations.py
+++ b/pypy/module/micronumpy/test/test_deprecations.py
@@ -24,7 +24,7 @@
                 # boolean indexing matches the dims in index
                 # to the first index.ndims in arr, not implemented in pypy yet
                 raises(IndexError, arr.__getitem__, index)
-                raises(TypeError, arr.__getitem__, (slice(None), index))
+                raises(IndexError, arr.__getitem__, (slice(None), index))
             else:
                 raises(np.VisibleDeprecationWarning, arr.__getitem__, index)
                 raises(np.VisibleDeprecationWarning, arr.__getitem__, 
(slice(None), index))
diff --git a/pypy/module/micronumpy/test/test_ndarray.py 
b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -2541,6 +2541,23 @@
         a[b] = np.array([[4.]])
         assert (a == [[4., 4., 4.]]).all()
 
+    def test_indexing_by_boolean(self):
+        import numpy as np
+        a = np.arange(6).reshape(2,3)
+        assert (a[[True, False], :] == [[3, 4, 5], [0, 1, 2]]).all()
+        b = a[np.array([True, False]), :]
+        assert (b == [[0, 1, 2]]).all()
+        assert b.base is None
+        b = a[:, np.array([True, False, True])]
+        assert b.base is not None
+
+    def test_scalar_indexing(self):
+        import numpy as np
+        a = np.arange(6).reshape(2,3)
+        i = np.dtype('int32').type(0)
+        assert (a[0] == a[i]).all()
+        
+
     def test_ellipsis_indexing(self):
         import numpy as np
         import sys
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to