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