Added arctan2 math function for numpy (ivy-llc#23315)

Co-authored: paulaehab<eng.paulaehab@gmail.com>
iababio · Sep 27, 2023 · d6c4b7e · d6c4b7e
1 parent 72190e6
commit d6c4b7e
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 0 deletions.
diff --git a/ivy/functional/frontends/numpy/__init__.py b/ivy/functional/frontends/numpy/__init__.py
@@ -565,6 +565,7 @@ def promote_types_of_numpy_inputs(
     _sin,
     _tan,
     _degrees,
+    _arctan2,
 )
 
 from ivy.functional.frontends.numpy.mathematical_functions.handling_complex_numbers import (  # noqa
@@ -672,6 +673,7 @@ def promote_types_of_numpy_inputs(
 arccos = ufunc("_arccos")
 arcsin = ufunc("_arcsin")
 arctan = ufunc("_arctan")
+arctan2 = ufunc("_arctan2")
 cos = ufunc("_cos")
 deg2rad = ufunc("_deg2rad")
 rad2deg = ufunc("_rad2deg")

diff --git a/ivy/functional/frontends/numpy/mathematical_functions/trigonometric_functions.py b/ivy/functional/frontends/numpy/mathematical_functions/trigonometric_functions.py
@@ -103,6 +103,32 @@ def _arctan(
     return ret
 
 
+# arctan2
+
+
+@handle_numpy_out
+@handle_numpy_dtype
+@to_ivy_arrays_and_back
+@handle_numpy_casting
+@from_zero_dim_arrays_to_scalar
+def _arctan2(
+    x1,
+    x2,
+    /,
+    out=None,
+    *,
+    where=True,
+    casting="same_kind",
+    order="K",
+    dtype=None,
+    subok=True,
+):
+    ret = ivy.atan2(x1, x2, out=out)
+    if ivy.is_array(where):
+        ret = ivy.where(where, ret, ivy.default(out, ivy.zeros_like(ret)), out=out)
+    return ret
+
+
 @handle_numpy_out
 @handle_numpy_dtype
 @to_ivy_arrays_and_back

diff --git a/...ivy/test_frontends/test_numpy/test_mathematical_functions/test_trigonometric_functions.py b/...ivy/test_frontends/test_numpy/test_mathematical_functions/test_trigonometric_functions.py
@@ -198,6 +198,57 @@ def test_numpy_arctan(
     )
 
 
+# Test for arctan2
+@handle_frontend_test(
+    fn_tree="numpy.arctan2",
+    dtypes_values_casting=np_frontend_helpers.dtypes_values_casting_dtype(
+        arr_func=[
+            lambda: helpers.dtype_and_values(
+                available_dtypes=helpers.get_dtypes("float"),
+                num_arrays=2,
+            )
+        ],
+    ),
+    where=np_frontend_helpers.where(),
+    number_positional_args=np_frontend_helpers.get_num_positional_args_ufunc(
+        fn_name="arctan2"
+    ),
+)
+def test_numpy_arctan2(
+    dtypes_values_casting,
+    where,
+    frontend,
+    test_flags,
+    fn_tree,
+    backend_fw,
+    on_device,
+):
+    input_dtypes, x, casting, dtype = dtypes_values_casting  # Unpack x
+
+    where, input_dtypes, test_flags = np_frontend_helpers.handle_where_and_array_bools(
+        where=where,
+        input_dtype=input_dtypes,
+        test_flags=test_flags,
+    )
+    np_frontend_helpers.test_frontend_function(
+        input_dtypes=input_dtypes,
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        atol=1e-3,
+        x1=x[0],  # Input x1
+        x2=x[1],  # Input x2
+        out=None,
+        where=where,
+        casting=casting,
+        order="K",
+        dtype=dtype,
+        subok=True,
+    )
+
+
 # cos
 @handle_frontend_test(
     fn_tree="numpy.cos",