sympy_compat: Add atan2.

Author: mattwala

symengine/lib/symengine.pxd

@@ -256,6 +256,7 @@ cdef extern from "<symengine/basic.h>" namespace "SymEngine":
     bool is_a_ComplexMPC "SymEngine::is_a<SymEngine::ComplexMPC>"(const Basic &b) nogil
     bool is_a_Log "SymEngine::is_a<SymEngine::Log>"(const Basic &b) nogil
     bool is_a_PyNumber "SymEngine::is_a<SymEngine::PyNumber>"(const Basic &b) nogil
+    bool is_a_ATan2 "SymEngine::is_a<SymEngine::ATan2>"(const Basic &b) nogil
 
     RCP[const Basic] expand(RCP[const Basic] &o) nogil except +
 
@@ -417,6 +418,7 @@ cdef extern from "<symengine/functions.h>" namespace "SymEngine":
     cdef RCP[const Basic] function_symbol(string name, const vec_basic &arg) nogil except+
     cdef RCP[const Basic] abs(RCP[const Basic] &arg) nogil except+
     cdef RCP[const Basic] gamma(RCP[const Basic] &arg) nogil except+
+    cdef RCP[const Basic] atan2(RCP[const Basic] &num, RCP[const Basic] &den) nogil except+
 
     cdef cppclass Function(Basic):
         pass
@@ -512,6 +514,9 @@ cdef extern from "<symengine/functions.h>" namespace "SymEngine":
     cdef cppclass Gamma(Function):
         pass
 
+    cdef cppclass ATan2(Function):
+        pass
+
 IF HAVE_SYMENGINE_MPFR:
     cdef extern from "mpfr.h":
         ctypedef struct __mpfr_struct:

symengine/lib/symengine_wrapper.pxd

@@ -121,6 +121,9 @@ cdef class Abs(Function):
 cdef class Gamma(Function):
     pass
 
+cdef class ATan2(Function):
+    pass
+
 cdef class Derivative(Basic):
     pass
 

symengine/lib/symengine_wrapper.pyx

@@ -99,6 +99,8 @@ cdef c2py(RCP[const symengine.Basic] o):
         r = ATanh.__new__(ATanh)
     elif (symengine.is_a_ACoth(deref(o))):
         r = ACoth.__new__(ACoth)
+    elif (symengine.is_a_ATan2(deref(o))):
+        r = ATan2.__new__(ATan2)
     elif (symengine.is_a_PyNumber(deref(o))):
         r = PyNumber.__new__(PyNumber)
     else:
@@ -169,6 +171,8 @@ def sympy2symengine(a, raise_error=False):
             return acsc(a.args[0])
         elif isinstance(a, sympy.asec):
             return asec(a.args[0])
+        elif isinstance(a, sympy.atan2):
+            return atan2(*a.args)
     elif isinstance(a, sympy.functions.elementary.hyperbolic.HyperbolicFunction):
         if isinstance(a, sympy.sinh):
             return sinh(a.args[0])
@@ -2159,6 +2163,11 @@ def gamma(x):
     cdef Basic X = sympify(x)
     return c2py(symengine.gamma(X.thisptr))
 
+def atan2(x, y):
+    cdef Basic X = sympify(x)
+    cdef Basic Y = sympify(y)
+    return c2py(symengine.atan2(X.thisptr, Y.thisptr))
+
 def eval_double(x):
     cdef Basic X = sympify(x)
     return c2py(<RCP[const symengine.Basic]>(symengine.real_double(symengine.eval_double(deref(X.thisptr)))))

symengine/sympy_compat.py

@@ -240,6 +240,12 @@ def __new__(cls, a):
         return symengine.acoth(a)
 
 
+class atan2(_RegisteredFunction):
+    _classes = (symengine.ATan2,)
+
+    def __new__(cls, a, b):
+        return symengine.atan2(a, b)
+
 '''
 for i in ("""Sin Cos Tan Gamma Cot Csc Sec ASin ACos ATan
           ACot ACsc ASec Sinh Cosh Tanh Coth ASinh ACosh ATanh

symengine/tests/test_sympy_compat.py

@@ -1,5 +1,5 @@
 from symengine.sympy_compat import (Integer, Rational, S, Basic, Add, Mul,
-    Pow, symbols, Symbol, log, sin, zeros)
+    Pow, symbols, Symbol, log, sin, zeros, atan2)
 
 def test_Integer():
     i = Integer(5)
@@ -50,6 +50,13 @@ def test_log():
     i = log(x)
     assert isinstance(i, log)
 
+def test_ATan2():
+    x, y = symbols("x y")
+    i = atan2(x, y)
+    assert isinstance(i, atan2)
+    i = atan2(0, 1)
+    assert i == 0
+
 def test_zeros():
     assert zeros(3, c=2).shape == (3, 2)