[pypy-commit] pypy default: Adds rint ufunc.

[taavi_burns]

Author: Taavi Burns <taavi.burns at gmail.com>
Branch: 
Changeset: r66117:8354267ccd65
Date: 2013-08-13 10:30 -0400
http://bitbucket.org/pypy/pypy/changeset/8354267ccd65/

Log:	Adds rint ufunc.

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -140,6 +140,7 @@
         ("deg2rad", "radians"),
         ("rad2deg", "degrees"),
         ("reciprocal", "reciprocal"),
+        ("rint", "rint"),
         ("sign", "sign"),
         ("signbit", "signbit"),
         ("sin", "sin"),
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -620,6 +620,7 @@
             ("positive", "pos", 1),
             ("negative", "neg", 1),
             ("absolute", "abs", 1, {"complex_to_float": True}),
+            ("rint", "rint", 1),
             ("sign", "sign", 1, {"promote_bools": True}),
             ("signbit", "signbit", 1, {"bool_result": True,
                                        "allow_complex": False}),
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -255,6 +255,22 @@
         for i in range(3):
             assert c[i] == a[i] * b[i]
 
+    def test_rint(self):
+        from numpypy import array, complex, rint, isnan
+
+        nnan, nan, inf, ninf = float('-nan'), float('nan'), float('inf'), float('-inf')
+
+        reference = array([ninf, -2., -1., -0., 0., 0., 0., 1., 2., inf])
+        a = array([ninf, -1.5, -1., -0.5, -0., 0., 0.5, 1., 1.5, inf])
+        b = rint(a)
+        for i in range(len(a)):
+            assert b[i] == reference[i]
+        assert isnan(rint(nan))
+        assert isnan(rint(nnan))
+
+        assert rint(complex(inf, 1.5)) == complex(inf, 2.)
+        assert rint(complex(0.5, inf)) == complex(0., inf)
+
     def test_sign(self):
         from numpypy import array, sign, dtype
 
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -307,6 +307,13 @@
     def min(self, v1, v2):
         return min(v1, v2)
 
+    @simple_unary_op
+    def rint(self, v):
+        if isfinite(v):
+            return rfloat.round_double(v, 0, half_even=True)
+        else:
+            return v
+
 class NonNativePrimitive(Primitive):
     _mixin_ = True
 
@@ -1334,6 +1341,15 @@
     #    return (rfloat.copysign(v1[0], v2[0]),
     #           rfloat.copysign(v1[1], v2[1]))
 
+    @specialize.argtype(1)
+    def rint(self, v):
+        ans = list(self.for_computation(self.unbox(v)))
+        if isfinite(ans[0]):
+            ans[0] = rfloat.round_double(ans[0], 0, half_even=True)
+        if isfinite(ans[1]):
+            ans[1] = rfloat.round_double(ans[1], 0, half_even=True)
+        return self.box_complex(ans[0], ans[1])
+
     @complex_unary_op
     def sign(self, v):
         '''