feat: adding support for numpy.real, imag, round, angle, (#3053)

* 2917: Adding ufunc-like functions real, imag, and angle.

These need unit tests still.

* Adding unit tests for real, imag, angle.

* More unit tests, fixing bugs

Fixing bugs in TypeTracer.angle.
Simplifying test_complex_ops.
Adding tests of TypeTracer real, imag, angle.

* Adding round function (almost numpy.ufunc)

* style: pre-commit fixes

* Fix dependence on numpy type conversions

As per Jim's recommendation, for real, imag, angle.

* Splitting real, imag, angle

* Adding to mostly-augen docs, removing round(.., out=...)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jim Pivarski <jpivarski@users.noreply.github.com>

docs/reference/toctree.txt

@@ -230,6 +230,10 @@
     generated/ak.nan_to_num
     generated/ak.values_astype
     generated/ak.strings_astype
+    generated/ak.round
+    generated/ak.real
+    generated/ak.imag
+    generated/ak.angle
 
 .. toctree::
     :caption: Arrays of categorical data

src/awkward/_nplikes/array_module.py

@@ -631,6 +631,26 @@ def cumsum(
         assert not isinstance(x, PlaceholderArray)
         return self._module.cumsum(x, axis=axis, out=maybe_out)
 
+    def real(self, x: ArrayLikeT) -> ArrayLikeT:
+        assert not isinstance(x, PlaceholderArray)
+        xr = self._module.real(x)
+        # For numpy, xr is a view on x, but we don't want to mutate x.
+        return self._module.copy(xr)
+
+    def imag(self, x: ArrayLikeT) -> ArrayLikeT:
+        assert not isinstance(x, PlaceholderArray)
+        xr = self._module.imag(x)
+        # For numpy, xr is a view on x, but we don't want to mutate x.
+        return self._module.copy(xr)
+
+    def angle(self, x: ArrayLikeT, deg: bool = False) -> ArrayLikeT:
+        assert not isinstance(x, PlaceholderArray)
+        return self._module.angle(x, deg)
+
+    def round(self, x: ArrayLikeT, decimals: int = 0) -> ArrayLikeT:
+        assert not isinstance(x, PlaceholderArray)
+        return self._module.round(x, decimals=decimals)
+
     def array_str(
         self,
         x: ArrayLikeT,

src/awkward/_nplikes/numpy_like.py

@@ -545,6 +545,22 @@ def can_cast(self, from_: DType | ArrayLikeT, to: DType | ArrayLikeT) -> bool:
     def is_c_contiguous(self, x: ArrayLikeT | PlaceholderArray) -> bool:
         ...
 
+    @abstractmethod
+    def real(self, x: ArrayLikeT) -> ArrayLikeT:
+        ...
+
+    @abstractmethod
+    def imag(self, x: ArrayLikeT) -> ArrayLikeT:
+        ...
+
+    @abstractmethod
+    def angle(self, x: ArrayLikeT, deg: bool = False) -> ArrayLikeT:
+        ...
+
+    @abstractmethod
+    def round(self, x: ArrayLikeT, decimals: int = 0) -> ArrayLikeT:
+        ...
+
     @classmethod
     @abstractmethod
     def is_own_array(cls, obj) -> bool:

src/awkward/_nplikes/typetracer.py

@@ -1438,6 +1438,33 @@ def isnan(self, x: TypeTracerArray) -> TypeTracerArray:
         try_touch_data(x)
         return TypeTracerArray._new(np.dtype(np.bool_), shape=x.shape)
 
+    def real(self, x: TypeTracerArray) -> TypeTracerArray:
+        assert isinstance(x, TypeTracerArray)
+        try_touch_data(x)
+        real_type = numpy.real(numpy.zeros(0, dtype=x.dtype)).dtype
+        return TypeTracerArray._new(real_type, shape=x.shape)
+
+    def imag(self, x: TypeTracerArray) -> TypeTracerArray:
+        assert isinstance(x, TypeTracerArray)
+        try_touch_data(x)
+        real_type = numpy.imag(numpy.zeros(0, dtype=x.dtype)).dtype
+        return TypeTracerArray._new(real_type, shape=x.shape)
+
+    def angle(self, x: TypeTracerArray, deg: bool = False) -> TypeTracerArray:
+        assert isinstance(x, TypeTracerArray)
+        try_touch_data(x)
+        float_type = numpy.angle(numpy.zeros(0, dtype=x.dtype)).dtype
+        return TypeTracerArray._new(float_type, shape=x.shape)
+
+    def round(
+        self,
+        x: TypeTracerArray,
+        decimals: int = 0,
+    ) -> TypeTracerArray:
+        assert isinstance(x, TypeTracerArray)
+        try_touch_data(x)
+        return TypeTracerArray._new(x.dtype, shape=x.shape)
+
     ############################ reducers
 
     def all(

src/awkward/operations/__init__.py

@@ -5,6 +5,7 @@
 import awkward.operations.str
 from awkward.operations.ak_all import *
 from awkward.operations.ak_almost_equal import *
+from awkward.operations.ak_angle import *
 from awkward.operations.ak_any import *
 from awkward.operations.ak_argcartesian import *
 from awkward.operations.ak_argcombinations import *
@@ -45,6 +46,7 @@
 from awkward.operations.ak_from_rdataframe import *
 from awkward.operations.ak_from_regular import *
 from awkward.operations.ak_full_like import *
+from awkward.operations.ak_imag import *
 from awkward.operations.ak_is_categorical import *
 from awkward.operations.ak_is_none import *
 from awkward.operations.ak_is_tuple import *
@@ -69,6 +71,8 @@
 from awkward.operations.ak_prod import *
 from awkward.operations.ak_ptp import *
 from awkward.operations.ak_ravel import *
+from awkward.operations.ak_real import *
+from awkward.operations.ak_round import *
 from awkward.operations.ak_run_lengths import *
 from awkward.operations.ak_singletons import *
 from awkward.operations.ak_softmax import *

src/awkward/operations/ak_angle.py

@@ -0,0 +1,55 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import awkward as ak
+from awkward._backends.numpy import NumpyBackend
+from awkward._dispatch import high_level_function
+from awkward._layout import HighLevelContext
+from awkward._nplikes.numpy_like import NumpyMetadata
+
+__all__ = ("angle",)
+
+np = NumpyMetadata.instance()
+cpu = NumpyBackend.instance()
+
+
+@ak._connect.numpy.implements("angle")
+@high_level_function()
+def angle(val, deg=False, highlevel=True, behavior=None, attrs=None):
+    """
+    Args:
+        val : array_like
+            Input array.
+        deg (bool, default is False): If True, returns angles in degrees,
+            otherwise in radians.
+        highlevel (bool, default is True): If True, return an #ak.Array;
+            otherwise, return a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
+        attrs (None or dict): Custom attributes for the output array, if
+            high-level.
+
+    Returns the counterclockwise angle from the positive real axis on the complex
+    plane in the range ``(-pi, pi]``, with dtype as a float.
+    """
+    # Dispatch
+    yield (val,)
+
+    # Implementation
+    return _impl_angle(val, deg, highlevel, behavior, attrs)
+
+
+def _impl_angle(val, deg, highlevel, behavior, attrs):
+    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
+        layout = ctx.unwrap(val, allow_record=False, primitive_policy="error")
+
+    # A closure over deg:
+    def action_angle(layout, backend, **kwargs):
+        if isinstance(layout, ak.contents.NumpyArray):
+            return ak.contents.NumpyArray(backend.nplike.angle(layout.data, deg))
+        else:
+            return None
+
+    out = ak._do.recursively_apply(layout, action_angle)
+    return ctx.wrap(out, highlevel=highlevel)

src/awkward/operations/ak_imag.py

@@ -0,0 +1,53 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import awkward as ak
+from awkward._backends.numpy import NumpyBackend
+from awkward._dispatch import high_level_function
+from awkward._layout import HighLevelContext
+from awkward._nplikes.numpy_like import NumpyMetadata
+
+__all__ = ("imag",)
+
+np = NumpyMetadata.instance()
+cpu = NumpyBackend.instance()
+
+
+@ak._connect.numpy.implements("imag")
+@high_level_function()
+def imag(val, highlevel=True, behavior=None, attrs=None):
+    """
+    Args:
+        val : array_like
+            Input array.
+        highlevel (bool, default is True): If True, return an #ak.Array;
+            otherwise, return a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
+        attrs (None or dict): Custom attributes for the output array, if
+            high-level.
+
+    Returns the imaginary components of the given array elements.
+    If the arrays have complex elements, the returned arrays are floats.
+    """
+    # Dispatch
+    yield (val,)
+
+    # Implementation
+    return _impl_imag(val, highlevel, behavior, attrs)
+
+
+def _impl_imag(val, highlevel, behavior, attrs):
+    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
+        layout = ctx.unwrap(val, allow_record=False, primitive_policy="error")
+
+    out = ak._do.recursively_apply(layout, _action_imag)
+    return ctx.wrap(out, highlevel=highlevel)
+
+
+def _action_imag(layout, backend, **kwargs):
+    if isinstance(layout, ak.contents.NumpyArray):
+        return ak.contents.NumpyArray(backend.nplike.imag(layout.data))
+    else:
+        return None

src/awkward/operations/ak_real.py

@@ -0,0 +1,53 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import awkward as ak
+from awkward._backends.numpy import NumpyBackend
+from awkward._dispatch import high_level_function
+from awkward._layout import HighLevelContext
+from awkward._nplikes.numpy_like import NumpyMetadata
+
+__all__ = ("real",)
+
+np = NumpyMetadata.instance()
+cpu = NumpyBackend.instance()
+
+
+@ak._connect.numpy.implements("real")
+@high_level_function()
+def real(val, highlevel=True, behavior=None, attrs=None):
+    """
+    Args:
+        val : array_like
+            Input array.
+        highlevel (bool, default is True): If True, return an #ak.Array;
+            otherwise, return a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
+        attrs (None or dict): Custom attributes for the output array, if
+            high-level.
+
+    Returns the real components of the given array elements.
+    If the arrays have complex elements, the returned arrays are floats.
+    """
+    # Dispatch
+    yield (val,)
+
+    # Implementation
+    return _impl_real(val, highlevel, behavior, attrs)
+
+
+def _impl_real(val, highlevel, behavior, attrs):
+    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
+        layout = ctx.unwrap(val, allow_record=False, primitive_policy="error")
+
+    out = ak._do.recursively_apply(layout, _action_real)
+    return ctx.wrap(out, highlevel=highlevel)
+
+
+def _action_real(layout, backend, **kwargs):
+    if isinstance(layout, ak.contents.NumpyArray):
+        return ak.contents.NumpyArray(backend.nplike.real(layout.data))
+    else:
+        return None

src/awkward/operations/ak_round.py

@@ -0,0 +1,64 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import awkward as ak
+from awkward._connect.numpy import UNSUPPORTED
+from awkward._dispatch import high_level_function
+from awkward._layout import HighLevelContext
+from awkward._nplikes.numpy_like import NumpyMetadata
+
+__all__ = ("round",)
+
+np = NumpyMetadata.instance()
+
+
+@ak._connect.numpy.implements("round")
+@high_level_function()
+def round(
+    array,
+    decimals: int = 0,
+    out=UNSUPPORTED,
+    highlevel=True,
+    behavior=None,
+    attrs=None,
+):
+    """
+    Args:
+        array : array_like
+            Input array.
+        decimals : int, optional
+            Number of decimal places to round to (default: 0).  If
+            decimals is negative, it specifies the number of positions to
+            the left of the decimal point.
+        out : unsupported optional argument
+        highlevel (bool, default is True): If True, return an #ak.Array;
+            otherwise, return a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
+        attrs (None or dict): Custom attributes for the output array, if
+            high-level.
+
+    Returns the real components of the given array elements.
+    If the arrays have complex elements, the returned arrays are floats.
+    """
+    # Dispatch
+    yield (array,)
+
+    # Implementation
+    return _impl(array, decimals, highlevel, behavior, attrs)
+
+
+def _impl(array, decimals, highlevel, behavior, attrs):
+    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
+        layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
+
+    # A closure over deg:
+    def action(layout, backend, **kwargs):
+        if isinstance(layout, ak.contents.NumpyArray):
+            return ak.contents.NumpyArray(backend.nplike.round(layout.data, decimals))
+        else:
+            return None
+
+    out = ak._do.recursively_apply(layout, action)
+    return ctx.wrap(out, highlevel=highlevel)