Commits

mattip committed 79d8f7d Merge

merge ndarray-round which added round() to ndarray

  • Participants
  • Parent commits 666a488, c8a81e7

Comments (0)

Files changed (7)

File pypy/doc/whatsnew-head.rst

 Fix a segfault in argsort when sorting by chunks on multidim numpypy arrays (mikefc)
 
 .. branch: dtype-isnative
+.. branch: ndarray-round

File pypy/module/micronumpy/interp_boxes.py

         w_values = space.newtuple([self])
         return convert_to_array(space, w_values)
 
+    @unwrap_spec(decimals=int)
+    def descr_round(self, space, decimals=0):
+        v = self.convert_to(self.get_dtype(space))
+        return self.get_dtype(space).itemtype.round(v, decimals)
+
 class W_BoolBox(W_GenericBox, PrimitiveBox):
     descr__new__, _get_dtype, descr_reduce = new_dtype_getter("bool")
 
     any = interp2app(W_GenericBox.descr_any),
     all = interp2app(W_GenericBox.descr_all),
     ravel = interp2app(W_GenericBox.descr_ravel),
+    round = interp2app(W_GenericBox.descr_round),
 )
 
 W_BoolBox.typedef = TypeDef("bool_", W_GenericBox.typedef,

File pypy/module/micronumpy/interp_numarray.py

         raise OperationError(space.w_NotImplementedError, space.wrap(
             "resize not implemented yet"))
 
-    def descr_round(self, space, w_decimals=0, w_out=None):
-        raise OperationError(space.w_NotImplementedError, space.wrap(
-            "round not implemented yet"))
+    @unwrap_spec(decimals=int)
+    def descr_round(self, space, decimals=0, w_out=None):
+        if space.is_none(w_out):
+            if self.get_dtype().is_bool_type():
+                #numpy promotes bool.round() to float16. Go figure.
+                w_out = W_NDimArray.from_shape(self.get_shape(),
+                       interp_dtype.get_dtype_cache(space).w_float16dtype)
+            else:
+                w_out = None
+        elif not isinstance(w_out, W_NDimArray):
+            raise OperationError(space.w_TypeError, space.wrap(
+                "return arrays must be of ArrayType"))
+        out = interp_dtype.dtype_agreement(space, [self], self.get_shape(),
+                                           w_out)
+        if out.get_dtype().is_bool_type() and self.get_dtype().is_bool_type():
+            calc_dtype = interp_dtype.get_dtype_cache(space).w_longdtype
+        else:
+            calc_dtype = out.get_dtype()
+
+        loop.round(space, self, calc_dtype, self.get_shape(), decimals, out)
+        return out
 
     def descr_searchsorted(self, space, w_v, w_side='left'):
         raise OperationError(space.w_NotImplementedError, space.wrap(
     byteswap = interp2app(W_NDimArray.descr_byteswap),
     choose   = interp2app(W_NDimArray.descr_choose),
     clip     = interp2app(W_NDimArray.descr_clip),
+    round    = interp2app(W_NDimArray.descr_round),
     data     = GetSetProperty(W_NDimArray.descr_get_data),
     diagonal = interp2app(W_NDimArray.descr_diagonal),
 

File pypy/module/micronumpy/loop.py

         iter = x_iter
     shapelen = len(shape)
     while not iter.done():
-        where_driver.jit_merge_point(shapelen=shapelen, dtype=dtype, 
+        where_driver.jit_merge_point(shapelen=shapelen, dtype=dtype,
                                         arr_dtype=arr_dtype)
         w_cond = arr_iter.getitem()
         if arr_dtype.itemtype.bool(w_cond):
     return out
 
 axis_reduce__driver = jit.JitDriver(name='numpy_axis_reduce',
-                                    greens=['shapelen', 
+                                    greens=['shapelen',
                                             'func', 'dtype',
                                             'identity'],
                                     reds='auto')
     arg_driver = jit.JitDriver(name='numpy_' + op_name,
                                greens = ['shapelen', 'dtype'],
                                reds = 'auto')
-    
+
     def argmin_argmax(arr):
         result = 0
         idx = 1
      result.shape == [3, 5, 2, 4]
      broadcast shape should be [3, 5, 2, 7, 4]
      result should skip dims 3 which is len(result_shape) - 1
-        (note that if right is 1d, result should 
+        (note that if right is 1d, result should
                   skip len(result_shape))
      left should skip 2, 4 which is a.ndims-1 + range(right.ndims)
           except where it==(right.ndims-2)
     righti = right.create_dot_iter(broadcast_shape, right_skip)
     while not outi.done():
         dot_driver.jit_merge_point(dtype=dtype)
-        lval = lefti.getitem().convert_to(dtype) 
-        rval = righti.getitem().convert_to(dtype) 
-        outval = outi.getitem().convert_to(dtype) 
+        lval = lefti.getitem().convert_to(dtype)
+        rval = righti.getitem().convert_to(dtype)
+        outval = outi.getitem().convert_to(dtype)
         v = dtype.itemtype.mul(lval, rval)
         value = dtype.itemtype.add(v, outval).convert_to(dtype)
         outi.setitem(value)
         setitem_filter_driver.jit_merge_point(shapelen=shapelen,
                                               index_dtype=index_dtype,
                                               arr_dtype=arr_dtype,
-                                             ) 
+                                             )
         if index_iter.getitem_bool():
             arr_iter.setitem(value_iter.getitem())
             value_iter.next()
         out_iter.next()
         min_iter.next()
 
+round_driver = jit.JitDriver(greens = ['shapelen', 'dtype'],
+                                    reds = 'auto')
+
+def round(space, arr, dtype, shape, decimals, out):
+    arr_iter = arr.create_iter(shape)
+    shapelen = len(shape)
+    out_iter = out.create_iter(shape)
+    while not arr_iter.done():
+        round_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
+        w_v = dtype.itemtype.round(arr_iter.getitem().convert_to(dtype),
+                     decimals)
+        out_iter.setitem(w_v)
+        arr_iter.next()
+        out_iter.next()
+
 diagonal_simple_driver = jit.JitDriver(greens = ['axis1', 'axis2'],
                                        reds = 'auto')
 
         out_iter.setitem(arr.getitem_index(space, indexes))
         iter.next()
         out_iter.next()
-       
+

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

 
         a = zeros(3)
         assert loads(dumps(sum(a))) == sum(a)
+
+    def test_round(self):
+        from numpypy import int32, float64, complex128, bool
+        i = int32(1337)
+        f = float64(13.37)
+        c = complex128(13 + 37.j)
+        b = bool(0)
+        assert i.round(decimals=-2) == 1300
+        assert i.round(decimals=1) == 1337
+        assert c.round() == c
+        assert f.round() == 13.
+        assert f.round(decimals=-1) == 10.
+        assert f.round(decimals=1) == 13.4
+        exc = raises(AttributeError, 'b.round()')
+        assert exc.value[0] == "'bool' object has no attribute 'round'"
+

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

 
         skip('sign of nan is non-determinant')
         assert (signbit([float('nan'), float('-nan'), -float('nan')]) ==
-            [False, True, True]).all()    
+            [False, True, True]).all()
 
     def test_reciprocal(self):
         from numpypy import array, reciprocal, complex64, complex128
         assert all([math.copysign(1, f(abs(float("nan")))) == 1 for f in floor, ceil, trunc])
         assert all([math.copysign(1, f(-abs(float("nan")))) == -1 for f in floor, ceil, trunc])
 
+    def test_round(self):
+        from numpypy import array, dtype
+        ninf, inf = float("-inf"), float("inf")
+        a = array([ninf, -1.4, -1.5, -1.0, 0.0, 1.0, 1.4, 0.5, inf])
+        assert ([ninf, -1.0, -2.0, -1.0, 0.0, 1.0, 1.0, 0.0, inf] == a.round()).all()
+        i = array([-1000, -100, -1, 0, 1, 111, 1111, 11111], dtype=int)
+        assert (i == i.round()).all()
+        assert (i.round(decimals=4) == i).all()
+        assert (i.round(decimals=-4) == [0, 0, 0, 0, 0, 0, 0, 10000]).all()
+        b = array([True, False], dtype=bool)
+        bround = b.round()
+        assert (bround == [1., 0.]).all()
+        assert bround.dtype is dtype('float16')
+        c = array([10.5+11.5j, -15.2-100.3456j, 0.2343+11.123456j])
+        assert (c.round(0) == [10.+12.j, -15-100j, 0+11j]).all()
+
+
     def test_copysign(self):
         from numpypy import array, copysign
 
             assert b[i] == res
 
     def test_exp2(self):
-        import math 
+        import math
         from numpypy import array, exp2
         inf = float('inf')
         ninf = -float('inf')
              complex(inf, inf), complex(inf, ninf), complex(0, inf),
              complex(ninf, ninf), complex(nan, 0), complex(0, nan),
              complex(nan, nan)]
-        assert (isfinite(a) == [True, True, False, False, False, 
-                        False, False, False, False, False]).all() 
+        assert (isfinite(a) == [True, True, False, False, False,
+                        False, False, False, False, False]).all()
 
     def test_logical_ops(self):
         from numpypy import logical_and, logical_or, logical_xor, logical_not
         #numpy returns (a.real*b.real + a.imag*b.imag) / abs(b)**2
         expect = [3., -23., 1.]
         for i in range(len(a)):
-            assert b[i] == expect[i] 
+            assert b[i] == expect[i]
         b = floor_divide(a[0], 0.)
         assert math.isnan(b.real)
         assert b.imag == 0.

File pypy/module/micronumpy/types.py

         return self.box(
             func(
                 self,
-                self.for_computation(raw)
+                self.for_computation(raw),
             )
         )
     return dispatcher
             return v
         return 0
 
+    @specialize.argtype(1)
+    def round(self, v, decimals=0):
+        raw = self.for_computation(self.unbox(v))
+        if decimals < 0:
+            # No ** in rpython
+            factor = 1
+            for i in xrange(-decimals):
+                factor *=10
+            #int does floor division, we want toward zero
+            if raw < 0:
+                ans = - (-raw / factor * factor)
+            else:
+                ans = raw / factor * factor
+        else:
+            ans = raw
+        return self.box(ans)
+
     @raw_unary_op
     def signbit(self, v):
         return v < 0
     def ceil(self, v):
         return math.ceil(v)
 
+    @specialize.argtype(1)
+    def round(self, v, decimals=0):
+        raw = self.for_computation(self.unbox(v))
+        if rfloat.isinf(raw):
+            return v
+        elif rfloat.isnan(raw):
+            return v
+        ans = rfloat.round_double(raw, decimals, half_even=True)
+        return self.box(ans)
+
     @simple_unary_op
     def trunc(self, v):
         if v < 0:
         except ZeroDivisionError:
             return rfloat.NAN, rfloat.NAN
 
+    @specialize.argtype(1)
+    def round(self, v, decimals=0):
+        ans = list(self.for_computation(self.unbox(v)))
+        if isfinite(ans[0]):
+            ans[0] = rfloat.round_double(ans[0], decimals,  half_even=True)
+        if isfinite(ans[1]):
+            ans[1] = rfloat.round_double(ans[1], decimals,  half_even=True)
+        return self.box_complex(ans[0], ans[1])
+
     # No floor, ceil, trunc in numpy for complex
     #@simple_unary_op
     #def floor(self, v):