ENH: Add angle function

Author: NeilGirdhar

docs/api-reference.md

@@ -6,6 +6,7 @@
     :nosignatures:
     :toctree: generated
 
+    angle
     apply_where
     argpartition
     at

src/array_api_extra/__init__.py

@@ -1,6 +1,7 @@
 """Extra array functions built on top of the array API standard."""
 
 from ._delegation import (
+    angle,
     argpartition,
     atleast_nd,
     cov,
@@ -32,6 +33,7 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "angle",
     "apply_where",
     "argpartition",
     "at",

src/array_api_extra/_delegation.py

@@ -19,6 +19,7 @@
 from ._lib._utils._typing import Array, DType
 
 __all__ = [
+    "angle",
     "atleast_nd",
     "cov",
     "create_diagonal",
@@ -32,6 +33,42 @@
 ]
 
 
+def angle(z: Array, /, *, deg: bool = False, xp: ModuleType | None = None) -> Array:
+    """
+    Return the angle of a complex argument.
+
+    Parameters
+    ----------
+    z : array
+        A complex-valued or real-valued array.
+    deg : bool, optional
+        Return angle in degrees if True, radians if False (default).
+    xp : array_namespace, optional
+        The standard-compatible namespace for `z`. Default: infer.
+
+    Returns
+    -------
+    array
+        The counterclockwise angle from the positive real axis on the complex
+        plane in the range ``(-pi, pi]``.
+
+    Notes
+    -----
+    This function passes the imaginary and real parts of the argument to
+    ``xp.atan2`` to compute the result.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), xp=xp)
+    Array([0.        , 1.57079633, 0.78539816], dtype=array_api_strict.float64)
+    """
+    if xp is None:
+        xp = array_namespace(z)
+    return _funcs.angle(z, deg=deg, xp=xp)
+
+
 def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
     """
     Recursively expand the dimension of an array to at least `ndim`.

src/array_api_extra/_lib/_funcs.py

@@ -23,6 +23,7 @@
 from ._utils._typing import Array, Device, DType
 
 __all__ = [
+    "angle",
     "apply_where",
     "atleast_nd",
     "broadcast_shapes",
@@ -38,6 +39,24 @@
 ]
 
 
+def angle(z: Array, /, *, deg: bool = False, xp: ModuleType) -> Array:
+    # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    if xp.isdtype(z.dtype, "complex floating"):
+        zimag = xp.imag(z)
+        zreal = xp.real(z)
+    else:
+        if not xp.isdtype(z.dtype, "real floating"):
+            z = xp.astype(z, default_dtype(xp, device=_compat.device(z)))
+        zimag = xp.zeros_like(z)
+        zreal = z
+
+    a = xp.atan2(zimag, zreal)
+    if deg:
+        a *= 180 / xp.pi
+    return a
+
+
 @overload
 def apply_where(  # numpydoc ignore=GL08
     cond: Array,
@@ -50,7 +69,6 @@ def apply_where(  # numpydoc ignore=GL08
     xp: ModuleType | None = None,
 ) -> Array: ...
 
-
 @overload
 def apply_where(  # numpydoc ignore=GL08
     cond: Array,

tests/test_funcs.py

@@ -13,6 +13,7 @@
 from typing_extensions import override
 
 from array_api_extra import (
+    angle,
     apply_where,
     argpartition,
     at,
@@ -49,6 +50,7 @@
 from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import lazy_xp_function
 
+lazy_xp_function(angle)
 lazy_xp_function(apply_where)
 lazy_xp_function(argpartition)
 lazy_xp_function(atleast_nd)
@@ -73,6 +75,59 @@
 lazy_xp_function(_funcs_searchsorted)
 
 
+class TestAngle:
+    def test_basic(self, xp: ModuleType):
+        x = xp.asarray(
+            [
+                1 + 3j,
+                np.sqrt(2) / 2.0 + 1j * np.sqrt(2) / 2,
+                1,
+                1j,
+                -1,
+                -1j,
+                1 - 3j,
+                -1 + 3j,
+            ],
+            dtype=xp.complex128,
+        )
+        expected = xp.asarray(
+            [
+                np.arctan(3.0 / 1.0),
+                np.arctan(1.0),
+                0,
+                np.pi / 2,
+                np.pi,
+                -np.pi / 2.0,
+                -np.arctan(3.0 / 1.0),
+                np.pi - np.arctan(3.0 / 1.0),
+            ],
+            dtype=xp.float64,
+        )
+
+        xp_assert_close(angle(x), expected, rtol=0, atol=1e-11)
+        xp_assert_close(angle(x, deg=True), expected * 180 / xp.pi, rtol=0, atol=1e-11)
+
+    def test_real(self, xp: ModuleType):
+        x = xp.asarray([0.0, -0.0, 1.0, -1.0])
+        expected = xp.asarray([0.0, xp.pi, 0.0, xp.pi], dtype=x.dtype)
+        xp_assert_close(angle(x), expected)
+
+    def test_integral(self, xp: ModuleType):
+        x = xp.asarray([0, -1, 1], dtype=xp.int32)
+        actual = angle(x)
+        expected = xp.asarray(
+            [0.0, xp.pi, 0.0], dtype=default_dtype(xp, device=get_device(x))
+        )
+        xp_assert_close(actual, expected)
+        assert actual.dtype == expected.dtype
+
+    @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp")
+    def test_xp(self, xp: ModuleType):
+        x = xp.asarray([1.0, 1.0j, 1 + 1j], dtype=xp.complex128)
+        expected = xp.asarray([0.0, xp.pi / 2, xp.pi / 4], dtype=xp.float64)
+        xp_assert_close(angle(x, xp=xp), expected)
+
+
 class TestApplyWhere:
     @staticmethod
     def f1(x: Array, y: Array | int = 10) -> Array: