diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 694bd15d..61ea792e 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -1,11 +1,11 @@ """Extra array functions built on top of the array API standard.""" from . import testing -from ._agnostic._elementwise import angle, apply_where +from ._agnostic._elementwise import apply_where from ._agnostic._inspection import default_dtype from ._at import at from ._creation import create_diagonal, one_hot -from ._elementwise import deg2rad, isclose, nan_to_num, rad2deg, sinc +from ._elementwise import angle, deg2rad, isclose, nan_to_num, rad2deg, sinc from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index from ._lazy import lazy_apply from ._linalg import kron diff --git a/src/array_api_extra/_agnostic/_elementwise.py b/src/array_api_extra/_agnostic/_elementwise.py index db36bf42..5f41426f 100644 --- a/src/array_api_extra/_agnostic/_elementwise.py +++ b/src/array_api_extra/_agnostic/_elementwise.py @@ -7,7 +7,6 @@ from .._at import at from .._lib import _compat, _helpers from .._lib._typing import Array, ArrayNamespace -from . import _inspection __all__ = [ "angle", @@ -291,46 +290,13 @@ def sinc(x: Array, /, *, xp: ArrayNamespace) -> Array: return xp.sin(y) / y -def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array: - """ - Return the angle of the complex argument. - - Parameters - ---------- - z : Array - Input 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 - ----- - Real input ``x`` is interpreted as ``x + 0j``. - - 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) - >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp) - Array([ 0., 90., 45.], dtype=array_api_strict.float64) - """ - if xp is None: - xp = _compat.array_namespace(z) +def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._elementwise`.""" 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, _inspection.default_dtype(xp, device=_compat.device(z))) zimag = xp.zeros_like(z) zreal = z a = xp.atan2(zimag, zreal) diff --git a/src/array_api_extra/_elementwise.py b/src/array_api_extra/_elementwise.py index b679191d..abc773ae 100644 --- a/src/array_api_extra/_elementwise.py +++ b/src/array_api_extra/_elementwise.py @@ -4,7 +4,59 @@ from ._lib import _compat, _helpers from ._lib._typing import Array, ArrayNamespace -__all__ = ["deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"] +__all__ = ["angle", "deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"] + + +def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array: + """ + Return the angle of the complex argument. + + Parameters + ---------- + z : array + Input array. Real input is interpreted as having zero imaginary part. + 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]``. + + 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) + >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp) + Array([ 0., 90., 45.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = _compat.array_namespace(z) + + if not xp.isdtype(z.dtype, ("real floating", "complex floating")): + z = xp.astype( + z, _agnostic._inspection.default_dtype(xp, device=_compat.device(z)) + ) + + if ( + _compat.is_numpy_namespace(xp) + or _compat.is_cupy_namespace(xp) + or _compat.is_dask_namespace(xp) + or _compat.is_jax_namespace(xp) + ): + return xp.angle(z, deg=deg) + + # Torch treats real negative zero as positive zero, unlike atan2(0, z). + if _compat.is_torch_namespace(xp) and xp.isdtype(z.dtype, "complex floating"): + result = xp.angle(z) + return result * 180 / xp.pi if deg else result + + return _agnostic._elementwise.angle(z, deg=deg, xp=xp) def deg2rad(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: diff --git a/tests/main/test_elementwise.py b/tests/main/test_elementwise.py index 4ae2973f..213dbd48 100644 --- a/tests/main/test_elementwise.py +++ b/tests/main/test_elementwise.py @@ -23,6 +23,7 @@ from array_api_extra._lib._typing import Array, ArrayNamespace, Device from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function +lazy_xp_function(angle) lazy_xp_function(apply_where) lazy_xp_function(deg2rad) lazy_xp_function(isclose)