[pypy-commit] pypy numpypy-complex2: test, implement complex floor_divide

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: numpypy-complex2
Changeset: r57214:8f9a0b7388a6
Date: 2012-09-07 13:21 +0300
http://bitbucket.org/pypy/pypy/changeset/8f9a0b7388a6/

Log:	test, implement complex floor_divide

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
@@ -297,6 +297,10 @@
         assert a[0] == 1
         assert a[1] == 0
 
+        a = sign(array([10+10j, -10+10j, 0+10j, 0-10j, 0+0j, 0-0j], dtype=complex))
+        ref = [1, -1, 1, -1, 0, 0]
+        assert (a == ref).all()
+
     def test_signbit(self):
         from _numpypy import signbit
 
@@ -879,10 +883,24 @@
 
     def test_floordiv(self):
         from _numpypy import floor_divide, array
+        import math
         a = array([1., 2., 3., 4., 5., 6., 6.01])
         b = floor_divide(a, 2.5)
         for i in range(len(a)):
             assert b[i] == a[i] // 2.5
+        
+        a = array([10+10j, -15-100j, 0+10j], dtype=complex)
+        b = floor_divide(a, 2.5)
+        for i in range(len(a)):
+            assert b[i] == a[i] // 2.5
+        b = floor_divide(a, 2.5+3j)
+        #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] 
+        b = floor_divide(a[0], 0.)
+        assert math.isnan(b.real)
+        assert b.imag == 0.
 
     def test_logaddexp(self):
         import math
@@ -1023,7 +1041,7 @@
             assert repr(abs(inf_c)) == 'inf'
             assert repr(abs(complex(float('nan'), float('nan')))) == 'nan'
 
-        assert False, 'untested: sign, floor_div, ' + \
+        assert False, 'untested: ' + \
                      'signbit, fabs, fmax, fmin, floor, ceil, trunc, ' + \
                      'exp2, expm1, isnan, isinf, isneginf, isposinf, ' + \
                      'isfinite, radians, degrees, log2, log1p, ' + \
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
@@ -1108,13 +1108,14 @@
         return min(v1, v2)
 
 
-    @simple_binary_op
+    @complex_binary_op
     def floordiv(self, v1, v2):
         try:
-            r, i = rcomplex.c_div(v1, v2)
-            return math.floor(r), 0
+            ab = v1[0]*v2[0] + v1[1]*v2[1]
+            bb = v2[0]*v2[0] + v2[1]*v2[1]
+            return math.floor(ab/bb), 0.
         except ZeroDivisionError:
-            return rfloat.NAN, 0
+            return rfloat.NAN, 0.
 
     #complex mod does not exist in numpy
     #@simple_binary_op
@@ -1130,11 +1131,21 @@
         return (rfloat.copysign(v1[0], v2[0]),
                rfloat.copysign(v1[1], v2[1]))
 
-    @simple_unary_op
+    @complex_unary_op
     def sign(self, v):
-        if v == 0.0:
-            return 0.0
-        return rfloat.copysign(1.0, v)
+        '''
+        sign of complex number could be either the point closest to the unit circle
+        or {-1,0,1}, for compatability with numpy we choose the latter
+        '''
+        if v[0] == 0.0:
+            if v[1] == 0:
+                return 0,0
+            if v[1] > 0:
+                return 1,0
+            return -1,0
+        if v[0] > 0:
+            return 1,0
+        return -1,0
 
     @raw_unary_op
     def signbit(self, v):