Commits

Romain Guillebert  committed eef3ab0 Merge

(mattip, rguillebert) Adds indexing by scalar, adds int conversion from scalar and single element array, fixes compress, indexing by an array with a smaller shape and the indexed object.

  • Participants
  • Parent commits 58a5238, 52ae1f1

Comments (0)

Files changed (8)

File pypy/module/micronumpy/arrayimpl/base.py

     def base(self):
         raise NotImplementedError
 
-    def create_iter(self, shape=None):
+    def create_iter(self, shape=None, backward_broadcast=False):
         raise NotImplementedError
 
 class BaseArrayIterator(object):

File pypy/module/micronumpy/arrayimpl/concrete.py

 
-from pypy.module.micronumpy.arrayimpl import base
+from pypy.module.micronumpy.arrayimpl import base, scalar
 from pypy.module.micronumpy import support, loop, iter
 from pypy.module.micronumpy.base import convert_to_array, W_NDimArray,\
      ArrayArgumentException
     parent = None
 
     # JIT hints that length of all those arrays is a constant
-    
+
     def get_shape(self):
         shape = self.shape
         jit.hint(len(shape), promote=True)
                               new_shape, self, orig_array)
         else:
             return None
-    
+
     def get_real(self, orig_array):
         strides = self.get_strides()
         backstrides = self.get_backstrides()
             dtype =  self.dtype.float_type
             return SliceArray(self.start, strides, backstrides,
                           self.get_shape(), self, orig_array, dtype=dtype)
-        return SliceArray(self.start, strides, backstrides, 
+        return SliceArray(self.start, strides, backstrides,
                           self.get_shape(), self, orig_array)
 
     def get_imag(self, orig_array):
         backstrides = self.get_backstrides()
         if self.dtype.is_complex_type():
             dtype =  self.dtype.float_type
-            return SliceArray(self.start + dtype.get_size(), strides, 
+            return SliceArray(self.start + dtype.get_size(), strides,
                     backstrides, self.get_shape(), self, orig_array, dtype=dtype)
         if self.dtype.is_flexible_type():
             # numpy returns self for self.imag
             space.isinstance_w(w_idx, space.w_slice) or
             space.is_w(w_idx, space.w_None)):
             raise IndexError
-        if isinstance(w_idx, W_NDimArray):
+        if isinstance(w_idx, W_NDimArray) and not isinstance(w_idx.implementation, scalar.Scalar):
             raise ArrayArgumentException
         shape = self.get_shape()
         shape_len = len(shape)
-        if shape_len == 0:
-            raise OperationError(space.w_IndexError, space.wrap(
-                "0-d arrays can't be indexed"))
         view_w = None
-        if (space.isinstance_w(w_idx, space.w_list) or
-            isinstance(w_idx, W_NDimArray)):
+        if space.isinstance_w(w_idx, space.w_list):
             raise ArrayArgumentException
         if space.isinstance_w(w_idx, space.w_tuple):
             view_w = space.fixedview(w_idx)
         shape = self.get_shape()[:]
         strides = self.get_strides()[:]
         backstrides = self.get_backstrides()[:]
-        shape[axis1], shape[axis2] = shape[axis2], shape[axis1]   
+        shape[axis1], shape[axis2] = shape[axis2], shape[axis1]
         strides[axis1], strides[axis2] = strides[axis2], strides[axis1]
-        backstrides[axis1], backstrides[axis2] = backstrides[axis2], backstrides[axis1] 
-        return W_NDimArray.new_slice(self.start, strides, 
+        backstrides[axis1], backstrides[axis2] = backstrides[axis2], backstrides[axis1]
+        return W_NDimArray.new_slice(self.start, strides,
                                      backstrides, shape, self, orig_arr)
 
     def get_storage_as_int(self, space):
         self.backstrides = backstrides
         self.storage = storage
 
-    def create_iter(self, shape=None):
+    def create_iter(self, shape=None, backward_broadcast=False):
         if shape is None or shape == self.get_shape():
             return iter.ConcreteArrayIterator(self)
         r = calculate_broadcast_strides(self.get_strides(),
                                         self.get_backstrides(),
-                                        self.get_shape(), shape)
+                                        self.get_shape(), shape, backward_broadcast)
         return iter.MultiDimViewIterator(self, self.dtype, 0, r[0], r[1], shape)
 
     def fill(self, box):
         free_raw_storage(self.storage, track_allocation=False)
 
 
-        
+
 
 class NonWritableArray(ConcreteArray):
     def descr_setitem(self, space, orig_array, w_index, w_value):
         raise OperationError(space.w_RuntimeError, space.wrap(
             "array is not writable"))
-        
+
 
 class SliceArray(BaseConcreteArray):
     def __init__(self, start, strides, backstrides, shape, parent, orig_arr,
     def fill(self, box):
         loop.fill(self, box.convert_to(self.dtype))
 
-    def create_iter(self, shape=None):
+    def create_iter(self, shape=None, backward_broadcast=False):
         if shape is not None and shape != self.get_shape():
             r = calculate_broadcast_strides(self.get_strides(),
                                             self.get_backstrides(),
-                                            self.get_shape(), shape)
+                                            self.get_shape(), shape,
+                                            backward_broadcast)
             return iter.MultiDimViewIterator(self.parent, self.dtype,
                                              self.start, r[0], r[1], shape)
         if len(self.get_shape()) == 1:
-            return iter.OneDimViewIterator(self.parent, self.dtype, self.start, 
+            return iter.OneDimViewIterator(self.parent, self.dtype, self.start,
                     self.get_strides(), self.get_shape())
         return iter.MultiDimViewIterator(self.parent, self.dtype, self.start,
                                     self.get_strides(),

File pypy/module/micronumpy/arrayimpl/scalar.py

     def get_strides(self):
         return []
 
-    def create_iter(self, shape=None):
+    def create_iter(self, shape=None, backward_broadcast=False):
         return ScalarIterator(self)
 
     def get_scalar_value(self):
     def descr_setitem(self, space, _, w_idx, w_val):
         raise OperationError(space.w_IndexError,
                              space.wrap("scalars cannot be indexed"))
-        
+
     def setitem_index(self, space, idx, w_val):
         raise OperationError(space.w_IndexError,
                              space.wrap("scalars cannot be indexed"))
 
     def reshape(self, space, orig_array, new_shape):
         return self.set_shape(space, orig_array, new_shape)
-        
+
     def create_axis_iter(self, shape, dim, cum):
         raise Exception("axis iter should not happen on scalar")
 

File pypy/module/micronumpy/interp_flatiter.py

     def get_shape(self):
         return self.shape
 
-    def create_iter(self, shape=None):
+    def create_iter(self, shape=None, backward_broadcast=False):
         assert isinstance(self.base(), W_NDimArray)
         return self.base().create_iter()
 

File pypy/module/micronumpy/interp_numarray.py

 from pypy.module.micronumpy import loop
 from pypy.module.micronumpy.dot import match_dot_shapes
 from pypy.module.micronumpy.interp_arrayops import repeat, choose
+from pypy.module.micronumpy.arrayimpl import scalar
 from rpython.tool.sourcetools import func_with_new_name
 from rpython.rlib import jit
 from rpython.rlib.rstring import StringBuilder
             raise OperationError(space.w_ValueError,
                                  space.wrap("index out of range for array"))
         size = loop.count_all_true(arr)
-        res = W_NDimArray.from_shape([size], self.get_dtype())
+        if len(arr.get_shape()) == 1:
+            res_shape = [size] + self.get_shape()[1:]
+        else:
+            res_shape = [size]
+        res = W_NDimArray.from_shape(res_shape, self.get_dtype())
         return loop.getitem_filter(res, self, arr)
 
     def setitem_filter(self, space, idx, val):
         s.append('])')
         return s.build()
 
-    def create_iter(self, shape=None):
+    def create_iter(self, shape=None, backward_broadcast=False):
         assert isinstance(self.implementation, BaseArrayImplementation)
-        return self.implementation.create_iter(shape)
+        return self.implementation.create_iter(shape=shape,
+                                   backward_broadcast=backward_broadcast)
 
     def create_axis_iter(self, shape, dim, cum):
         return self.implementation.create_axis_iter(shape, dim, cum)
         if not space.is_none(w_axis):
             raise OperationError(space.w_NotImplementedError,
                                  space.wrap("axis unsupported for compress"))
+            arr = self
+        else:
+            arr = self.descr_reshape(space, [space.wrap(-1)])
         index = convert_to_array(space, w_obj)
-        return self.getitem_filter(space, index)
+        return arr.getitem_filter(space, index)
 
     def descr_flatten(self, space, w_order=None):
         if self.is_scalar():
     descr_argmax = _reduce_argmax_argmin_impl("max")
     descr_argmin = _reduce_argmax_argmin_impl("min")
 
+    def descr_int(self, space):
+        shape = self.get_shape()
+        if len(shape) == 0:
+            assert isinstance(self.implementation, scalar.Scalar)
+            return space.int(space.wrap(self.implementation.get_scalar_value()))
+        if shape == [1]:
+            return space.int(self.descr_getitem(space, space.wrap(0)))
+        raise OperationError(space.w_TypeError, space.wrap("only length-1 arrays can be converted to Python scalars"))
+
 
 @unwrap_spec(offset=int)
 def descr_new_array(space, w_subtype, w_shape, w_dtype=None, w_buffer=None,
 
     __repr__ = interp2app(W_NDimArray.descr_repr),
     __str__ = interp2app(W_NDimArray.descr_str),
+    __int__ = interp2app(W_NDimArray.descr_int),
 
     __pos__ = interp2app(W_NDimArray.descr_pos),
     __neg__ = interp2app(W_NDimArray.descr_neg),

File pypy/module/micronumpy/loop.py

 
 def getitem_filter(res, arr, index):
     res_iter = res.create_iter()
-    index_iter = index.create_iter()
+    shapelen = len(arr.get_shape())
+    if shapelen > 1 and len(index.get_shape()) < 2:
+        index_iter = index.create_iter(arr.get_shape(), backward_broadcast=True)
+    else:
+        index_iter = index.create_iter()
     arr_iter = arr.create_iter()
-    shapelen = len(arr.get_shape())
     arr_dtype = arr.get_dtype()
     index_dtype = index.get_dtype()
     # XXX length of shape of index as well?

File pypy/module/micronumpy/strides.py

     rshape += shape[s:]
     return rshape, rstart, rstrides, rbackstrides
 
-def calculate_broadcast_strides(strides, backstrides, orig_shape, res_shape):
+def calculate_broadcast_strides(strides, backstrides, orig_shape, res_shape, backwards=False):
     rstrides = []
     rbackstrides = []
     for i in range(len(orig_shape)):
         else:
             rstrides.append(strides[i])
             rbackstrides.append(backstrides[i])
-    rstrides = [0] * (len(res_shape) - len(orig_shape)) + rstrides
-    rbackstrides = [0] * (len(res_shape) - len(orig_shape)) + rbackstrides
+    if backwards:
+        rstrides = rstrides + [0] * (len(res_shape) - len(orig_shape))  
+        rbackstrides = rbackstrides + [0] * (len(res_shape) - len(orig_shape)) 
+    else:
+        rstrides = [0] * (len(res_shape) - len(orig_shape)) + rstrides
+        rbackstrides = [0] * (len(res_shape) - len(orig_shape)) + rbackstrides
     return rstrides, rbackstrides
 
 def is_single_elem(space, w_elem, is_rec_type):

File pypy/module/micronumpy/test/test_numarray.py

         assert a[1] == 'xyz'
         assert a.imag[0] == 'abc'
         raises(TypeError, 'a.imag = "qop"')
-        a=array([[1+1j, 2-3j, 4+5j],[-6+7j, 8-9j, -2-1j]]) 
+        a=array([[1+1j, 2-3j, 4+5j],[-6+7j, 8-9j, -2-1j]])
         assert a.real[0,1] == 2
         a.real[0,1] = -20
         assert a[0,1].real == -20
         assert a[1,2].imag == 30
         a.real = 13
         assert a[1,1].real == 13
-        a=array([1+1j, 2-3j, 4+5j, -6+7j, 8-9j, -2-1j]) 
+        a=array([1+1j, 2-3j, 4+5j, -6+7j, 8-9j, -2-1j])
         a.real = 13
         assert a[3].real == 13
         a.imag = -5
         from numpypy import array
         # testcases from numpy docstring
         x = array([[1, 2, 3]])
-        assert (x.swapaxes(0, 1) == array([[1], [2], [3]])).all() 
+        assert (x.swapaxes(0, 1) == array([[1], [2], [3]])).all()
         x = array([[[0,1],[2,3]],[[4,5],[6,7]]]) # shape = (2, 2, 2)
-        assert (x.swapaxes(0, 2) == array([[[0, 4], [2, 6]], 
-                                           [[1, 5], [3, 7]]])).all() 
-        assert (x.swapaxes(0, 1) == array([[[0, 1], [4, 5]], 
+        assert (x.swapaxes(0, 2) == array([[[0, 4], [2, 6]],
+                                           [[1, 5], [3, 7]]])).all()
+        assert (x.swapaxes(0, 1) == array([[[0, 1], [4, 5]],
                                            [[2, 3], [6, 7]]])).all()
-        assert (x.swapaxes(1, 2) == array([[[0, 2], [1, 3]], 
+        assert (x.swapaxes(1, 2) == array([[[0, 2], [1, 3]],
                                            [[4, 6],[5, 7]]])).all()
 
         # more complex shape i.e. (2, 2, 3)
-        x = array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]) 
-        assert (x.swapaxes(0, 1) == array([[[1, 2, 3], [7, 8, 9]], 
-                                           [[4, 5, 6], [10, 11, 12]]])).all() 
-        assert (x.swapaxes(0, 2) == array([[[1, 7], [4, 10]], [[2, 8], [5, 11]], 
-                                           [[3, 9], [6, 12]]])).all() 
-        assert (x.swapaxes(1, 2) == array([[[1, 4], [2, 5], [3, 6]], 
-                                           [[7, 10], [8, 11],[9, 12]]])).all() 
+        x = array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])
+        assert (x.swapaxes(0, 1) == array([[[1, 2, 3], [7, 8, 9]],
+                                           [[4, 5, 6], [10, 11, 12]]])).all()
+        assert (x.swapaxes(0, 2) == array([[[1, 7], [4, 10]], [[2, 8], [5, 11]],
+                                           [[3, 9], [6, 12]]])).all()
+        assert (x.swapaxes(1, 2) == array([[[1, 4], [2, 5], [3, 6]],
+                                           [[7, 10], [8, 11],[9, 12]]])).all()
 
         # test slice
-        assert (x[0:1,0:2].swapaxes(0,2) == array([[[1], [4]], [[2], [5]], 
+        assert (x[0:1,0:2].swapaxes(0,2) == array([[[1], [4]], [[2], [5]],
                                                    [[3], [6]]])).all()
         # test virtual
-        assert ((x + x).swapaxes(0,1) == array([[[ 2,  4,  6], [14, 16, 18]], 
+        assert ((x + x).swapaxes(0,1) == array([[[ 2,  4,  6], [14, 16, 18]],
                                          [[ 8, 10, 12], [20, 22, 24]]])).all()
         assert array(1).swapaxes(10, 12) == 1
 
         assert (zeros(1)[[]] == []).all()
 
     def test_int_array_index_setitem(self):
-        from numpypy import array, arange, zeros
+        from numpypy import arange, zeros, array
         a = arange(10)
         a[[3, 2, 1, 5]] = zeros(4, dtype=int)
         assert (a == [0, 0, 0, 0, 4, 0, 6, 7, 8, 9]).all()
         assert (b[array([True, False, True])] == [0, 2]).all()
         raises(ValueError, "array([1, 2])[array([True, True, True])]")
         raises(ValueError, "b[array([[True, False], [True, False]])]")
+        a = array([[1,2,3],[4,5,6],[7,8,9]],int)
+        c = array([True,False,True],bool)
+        b = a[c]
+        assert (a[c] == [[1, 2, 3], [7, 8, 9]]).all()
 
     def test_bool_array_index_setitem(self):
         from numpypy import arange, array
         b = arange(5)
         b[array([True, False, True])] = [20, 21, 0, 0, 0, 0, 0]
-        assert (b == [20, 1, 21, 3, 4]).all() 
+        assert (b == [20, 1, 21, 3, 4]).all()
         raises(ValueError, "array([1, 2])[array([True, False, True])] = [1, 2, 3]")
 
     def test_weakref(self):
         b = array([1, 2, 3, 4])
         assert (a == b) == False
 
+    def test__int__(self):
+        from numpypy import array
+        assert int(array(1)) == 1
+        assert int(array([1])) == 1
+        assert raises(TypeError, "int(array([1, 2]))")
+        assert int(array([1.5])) == 1
+
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):
         assert isinstance(i['data'][0], int)
 
     def test_array_indexing_one_elem(self):
-        skip("not yet")
         from numpypy import array, arange
         raises(IndexError, 'arange(3)[array([3.5])]')
         a = arange(3)[array([1])]
         a = arange(10)
         assert (a.compress([True, False, True]) == [0, 2]).all()
         assert (a.compress([1, 0, 13]) == [0, 2]).all()
+        assert (a.compress([1, 0, 13]) == [0, 2]).all()
         assert (a.compress([1, 0, 13.5]) == [0, 2]).all()
         assert (a.compress(array([1, 0, 13.5], dtype='>f4')) == [0, 2]).all()
         assert (a.compress(array([1, 0, 13.5], dtype='<f4')) == [0, 2]).all()
         BaseNumpyAppTest.setup_class.im_func(cls)
         cls.w_data = cls.space.wrap(struct.pack('dddd', 1, 2, 3, 4))
         cls.w_fdata = cls.space.wrap(struct.pack('f', 2.3))
-        cls.w_float16val = cls.space.wrap('\x00E') # 5.0 in float16 
+        cls.w_float16val = cls.space.wrap('\x00E') # 5.0 in float16
         cls.w_float32val = cls.space.wrap(struct.pack('f', 5.2))
         cls.w_float64val = cls.space.wrap(struct.pack('d', 300.4))
         cls.w_ulongval = cls.space.wrap(struct.pack('L', 12))
         from numpypy import array, arange
         assert array(2.0).argsort() == 0
         nnp = self.non_native_prefix
-        for dtype in ['int', 'float', 'int16', 'float32', 'uint64', 
+        for dtype in ['int', 'float', 'int16', 'float32', 'uint64',
                         nnp + 'i2', complex]:
             a = array([6, 4, -1, 3, 8, 3, 256+20, 100, 101], dtype=dtype)
             c = a.copy()
             assert (a == c).all() # not modified
             a = arange(100)
             assert (a.argsort() == a).all()
-        raises(NotImplementedError, 'arange(10,dtype="float16").argsort()')    
+        raises(NotImplementedError, 'arange(10,dtype="float16").argsort()')
 
     def test_argsort_nd(self):
         from numpypy import array
 
     def test_argsort_axis(self):
         from numpypy import array
-        a = array([[4, 2], [1, 3]]) 
+        a = array([[4, 2], [1, 3]])
         assert (a.argsort(axis=None) == [2, 1, 3, 0]).all()
         assert (a.argsort(axis=-1) == [[1, 0], [0, 1]]).all()
         assert (a.argsort(axis=0) == [[1, 0], [0, 1]]).all()