Repository: scikit-hep/awkward

Issue:
np.real, np.imag, np.round, np.around can't take an Awkward Array as input
### Version of Awkward Array

HEAD

### Description and code to reproduce

This is more like a strange "gotcha" of NumPy, but `np.real` and `np.imag` are not ufuncs, unlike, for example, np.conjugate.

```python
>>> np.real
<function real at 0x7e73131d8930>
>>> np.imag
<function imag at 0x7e73131d8b30>
>>> np.conjugate
<ufunc 'conjugate'>
```

Therefore, it would require some special handling to make it work in Awkward, but not much. Here's what happens now:

```python
>>> np.real(ak.Array([[1+0.1j, 2+0.2j, 3+0.3j], [], [4+0.4j, 5+0.5j]]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/irishep/awkward/src/awkward/highlevel.py", line 1527, in __array_function__
    return ak._connect.numpy.array_function(
  File "/home/jpivarski/irishep/awkward/src/awkward/_connect/numpy.py", line 109, in array_function
    rectilinear_args = tuple(_to_rectilinear(x, backend) for x in args)
  File "/home/jpivarski/irishep/awkward/src/awkward/_connect/numpy.py", line 109, in <genexpr>
    rectilinear_args = tuple(_to_rectilinear(x, backend) for x in args)
  File "/home/jpivarski/irishep/awkward/src/awkward/_connect/numpy.py", line 78, in _to_rectilinear
    return layout.to_backend(backend).to_backend_array(allow_missing=True)
  File "/home/jpivarski/irishep/awkward/src/awkward/contents/content.py", line 1021, in to_backend_array
    return self._to_backend_array(allow_missing, backend)
  File "/home/jpivarski/irishep/awkward/src/awkward/contents/listoffsetarray.py", line 2078, in _to_backend_array
    return self.to_RegularArray()._to_backend_array(allow_missing, backend)
  File "/home/jpivarski/irishep/awkward/src/awkward/contents/listoffsetarray.py", line 288, in to_RegularArray
    self._backend.maybe_kernel_error(
  File "/home/jpivarski/irishep/awkward/src/awkward/_backends/backend.py", line 67, in maybe_kernel_error
    raise ValueError(self.format_kernel_error(error))
ValueError: cannot convert to RegularArray because subarray lengths are not regular (in compiled code: https://github.com/scikit-hep/awkward/blob/awkward-cpp-26/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_toRegularArray.cpp#L22)
```

It's trying to auto-convert the Awkward Array into a NumPy array to use the function, and it can't because this one is ragged.

Discussion:
tcawlfield: Are there automatic code formatters and formatting checks? I guess I'll find out! Oh, I did find out and there are indeed.
ianna: > Are there automatic code formatters and formatting checks? I guess I'll find out! Oh, I did find out and there are indeed.

Indeed, they are. You can run `pre-commit run --all` and fix the warnings before pushing then you will not need to pull the automatically updated branch from remote.
jpivarski: Let me know when you want me to re-review this!
jpivarski: Great, thanks! This one is done and I'll merge it.

Patch:
diff --git a/docs/reference/toctree.txt b/docs/reference/toctree.txt
index 9e025f47..8f004945 100644
--- a/docs/reference/toctree.txt
+++ b/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
diff --git a/src/awkward/_nplikes/array_module.py b/src/awkward/_nplikes/array_module.py
index 82ef7a11..916fe708 100644
--- a/src/awkward/_nplikes/array_module.py
+++ b/src/awkward/_nplikes/array_module.py
@@ -631,6 +631,26 @@ class ArrayModuleNumpyLike(NumpyLike[ArrayLikeT]):
         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,
diff --git a/src/awkward/_nplikes/numpy_like.py b/src/awkward/_nplikes/numpy_like.py
index 19119d75..949f32d2 100644
--- a/src/awkward/_nplikes/numpy_like.py
+++ b/src/awkward/_nplikes/numpy_like.py
@@ -545,6 +545,22 @@ class NumpyLike(PublicSingleton, Protocol[ArrayLikeT]):
     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:
diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py
index 1e2812be..28b00f30 100644
--- a/src/awkward/_nplikes/typetracer.py
+++ b/src/awkward/_nplikes/typetracer.py
@@ -1438,6 +1438,33 @@ class TypeTracer(NumpyLike[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(
diff --git a/src/awkward/operations/__init__.py b/src/awkward/operations/__init__.py
index 133ea8a5..90d4f991 100644
--- a/src/awkward/operations/__init__.py
+++ b/src/awkward/operations/__init__.py
@@ -5,6 +5,7 @@ from __future__ import annotations
 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_parquet import *
 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_parameters import *
 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 *
diff --git a/src/awkward/operations/ak_angle.py b/src/awkward/operations/ak_angle.py
new file mode 100644
index 00000000..2bf34fa3
--- /dev/null
+++ b/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)
diff --git a/src/awkward/operations/ak_imag.py b/src/awkward/operations/ak_imag.py
new file mode 100644
index 00000000..0bc4836e
--- /dev/null
+++ b/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
diff --git a/src/awkward/operations/ak_real.py b/src/awkward/operations/ak_real.py
new file mode 100644
index 00000000..2c49c98c
--- /dev/null
+++ b/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
diff --git a/src/awkward/operations/ak_round.py b/src/awkward/operations/ak_round.py
new file mode 100644
index 00000000..f0e6d3fe
--- /dev/null
+++ b/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)


Test patch:
diff --git a/tests/test_2917_real_imag_angle.py b/tests/test_2917_real_imag_angle.py
new file mode 100644
index 00000000..7246ad4e
--- /dev/null
+++ b/tests/test_2917_real_imag_angle.py
@@ -0,0 +1,83 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import numpy as np
+
+import awkward as ak
+from awkward.types import ListType
+from awkward.types.numpytype import NumpyType
+
+
+def test_complex_ops():
+    """
+    Testing Awkward implementations (cpu backend) of real, imag, and angle
+    """
+    arr_1 = ak.Array([[1 + 0.1j, 2 + 0.2j, 3 + 0.3j], [], [4 + 0.4j, 5 + 0.5j]])
+    arr_csingle = ak.from_numpy(np.array([6.0 + 0.0j, 0.0 + 7.0j], dtype="complex64"))
+    arr_real = ak.Array([11.0, 12.0])
+
+    real_1 = np.real(arr_1)
+    assert ak.all(real_1 == ak.Array([[1, 2, 3], [], [4, 5]]))
+
+    real_csingle = np.real(arr_csingle)
+    assert ak.all(real_csingle == ak.Array([6.0, 0.0]))
+    assert real_csingle.type.content == NumpyType("float32")
+
+    real_real = np.real(arr_real)
+    assert ak.all(real_real == ak.Array([11.0, 12.0]))
+
+    imag_1 = np.imag(arr_1)
+    assert ak.all(imag_1 == ak.Array([[0.1, 0.2, 0.3], [], [0.4, 0.5]]))
+
+    imag_csingle = np.imag(arr_csingle)
+    assert ak.all(imag_csingle == ak.Array([0.0, 7.0]))
+    assert imag_csingle.type.content == NumpyType("float32")
+
+    imag_real = np.imag(arr_real)
+    assert ak.all(imag_real == ak.Array([0.0, 0.0]))
+
+    angle_1 = np.angle(arr_1)
+    a1 = np.arctan(0.1)
+    expected = ak.Array([[a1, a1, a1], [], [a1, a1]])  # or [[a1] * 3, ...]
+    assert ak.all(np.abs(angle_1 - expected) < np.finfo(np.float64).eps)
+
+    angle_csingle = np.angle(arr_csingle, deg=True)
+    assert ak.all(angle_csingle == ak.Array([0.0, 90.0]))
+    assert angle_csingle.type.content == NumpyType("float32")
+
+    angle_real = np.angle(arr_real)
+    assert ak.all(angle_real == ak.Array([0.0, 0.0]))
+
+
+def test_complex_typetracer():
+    tt_arr = ak.to_backend(
+        ak.Array([[1 + 0.1j, 2 + 0.2j, 3 + 0.3j], [], [4 + 0.4j, 5 + 0.5j]]),
+        "typetracer",
+    )
+    tt_csingle = ak.to_backend(
+        ak.from_numpy(np.array([6.0 + 0.0j, 0.0 + 7.0j], dtype="complex64")),
+        "typetracer",
+    )
+    tt_real = ak.to_backend(
+        ak.from_numpy(np.array([11.0, 12.0], dtype="float16")), "typetracer"
+    )
+    tt_int = ak.to_backend(ak.Array([1, 2]), "typetracer")
+
+    real_arr = np.real(tt_arr)
+    assert real_arr.type.content == ListType(NumpyType("float64"))
+    real_tt_csingle = np.real(tt_csingle)
+    assert real_tt_csingle.type.content == NumpyType("float32")
+    real_tt_real = np.real(tt_real)
+    assert real_tt_real.type.content == NumpyType("float16")
+
+    imag_tt_csingle = np.imag(tt_csingle)
+    assert imag_tt_csingle.type.content == NumpyType("float32")
+    imag_tt_real = np.imag(tt_real)
+    assert imag_tt_real.type.content == NumpyType("float16")
+
+    angle_tt_csingle = np.angle(tt_csingle, deg=True)
+    assert angle_tt_csingle.type.content == NumpyType("float32")
+    angle_tt_real = np.angle(tt_real)
+    assert angle_tt_real.type.content == NumpyType("float16")
+    assert np.angle(tt_int).type.content == NumpyType("float64")
diff --git a/tests/test_2917_round.py b/tests/test_2917_round.py
new file mode 100644
index 00000000..f0793b7e
--- /dev/null
+++ b/tests/test_2917_round.py
@@ -0,0 +1,43 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import numpy as np
+
+import awkward as ak
+from awkward.operations import to_list
+
+
+def test_round():
+    """
+    Testing Awkward implementations (cpu backend) of real, imag, and angle
+    """
+    arr_1 = ak.Array(
+        [[1 + 0.111j, 2 + 0.222j, 3 + 0.333j], [], [4 + 0.444j, 5 + 0.555j]]
+    )
+    arr_2 = ak.Array([-2.5, -1.5, -0.6, -0.5, -0.4, 0.4, 0.5, 0.6, 1.5, 2.5])
+
+    assert to_list(np.round(arr_1)) == [
+        [(1 + 0j), (2 + 0j), (3 + 0j)],
+        [],
+        [(4 + 0j), (5 + 1j)],
+    ]
+
+    assert to_list(np.round(arr_1, decimals=2)) == [
+        [(1 + 0.11j), (2 + 0.22j), (3 + 0.33j)],
+        [],
+        [(4 + 0.44j), (5 + 0.56j)],
+    ]
+
+    assert to_list(np.round(arr_2)) == [
+        -2.0,
+        -2.0,
+        -1.0,
+        -0.0,
+        -0.0,
+        0.0,
+        0.0,
+        1.0,
+        2.0,
+        2.0,
+    ]