[pypy-commit] pypy numpy-ufuncs2: Adding arctan2

[taavi_burns]

Author: Taavi Burns <taavi.burns at gmail.com>
Branch: numpy-ufuncs2
Changeset: r53457:e90378882863
Date: 2012-03-13 11:15 -0700
http://bitbucket.org/pypy/pypy/changeset/e90378882863/

Log:	Adding arctan2

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
@@ -67,6 +67,7 @@
         ("arccos", "arccos"),
         ("arcsin", "arcsin"),
         ("arctan", "arctan"),
+        ("arctan2", "arctan2"),
         ("arccosh", "arccosh"),
         ("arcsinh", "arcsinh"),
         ("arctanh", "arctanh"),
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
@@ -440,6 +440,7 @@
             ("arcsin", "arcsin", 1, {"promote_to_float": True}),
             ("arccos", "arccos", 1, {"promote_to_float": True}),
             ("arctan", "arctan", 1, {"promote_to_float": True}),
+            ("arctan2", "arctan2", 2, {"promote_to_float": True}),
             ("sinh", "sinh", 1, {"promote_to_float": True}),
             ("cosh", "cosh", 1, {"promote_to_float": True}),
             ("tanh", "tanh", 1, {"promote_to_float": True}),
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
@@ -343,6 +343,21 @@
         b = arctan(a)
         assert math.isnan(b[0])
 
+    def test_arctan2(self):
+        import math
+        from _numpypy import array, arctan2
+
+        # From the numpy documentation
+        assert (
+            arctan2(
+                [0.,  0.,           1.,          -1., float('inf'),  float('inf')],
+                [0., -0., float('inf'), float('inf'), float('inf'), float('-inf')]) ==
+            [0.,  math.pi,  0., -0.,  math.pi/4, 3*math.pi/4]).all()
+
+        a = array([float('nan')])
+        b = arctan2(a, 0)
+        assert math.isnan(b[0])
+
     def test_sinh(self):
         import math
         from _numpypy import array, sinh
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
@@ -537,6 +537,10 @@
     def arctan(self, v):
         return math.atan(v)
 
+    @simple_binary_op
+    def arctan2(self, v1, v2):
+        return math.atan2(v1, v2)
+
     @simple_unary_op
     def sinh(self, v):
         return math.sinh(v)