VectorNumpy3D's azimuthal and longitudinal properties throw an error (similar for VectorNumpy4D)

[Saransh-cpp]

Reproducible example

import vector

vec = vector.array(
    [
        (1.1, 2.1, 3.1),
        (1.2, 2.2, 3.2),
        (1.3, 2.3, 3.3),
        (1.4, 2.4, 4.4),
        (1.5, 2.5, 5.5)
    ], dtype=[("x", float), ("y", float), ("z", float)]
)
print(vec.azimuthal)

Similarly for -

• vec.longitudinal
• vec.azimuthal for a 4D NumPy vector (VectorNumpy4D)
• vec.longitudinal for a 4D NumPy vector (VectorNumpy4D)
• vec.temporal for a 4D NumPy vector (VectorNumpy4D)

Error

>>> vec.azimuthal
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 1488, in _array_repr_implementation
    lst = array2string(arr, max_line_width, precision, suppress_small,
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 736, in array2string
    return _array2string(a, options, separator, prefix)
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 513, in wrapper
    return f(self, *args, **kwargs)
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 546, in _array2string
    lst = _formatArray(a, format_function, options['linewidth'],
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 889, in _formatArray
    return recurser(index=(),
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 853, in recurser
    word = recurser(index + (-1,), next_hanging_indent, next_width)
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\numpy\core\arrayprint.py", line 799, in recurser
    return format_function(a[index])
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\vector\_backends\numpy_.py", line 218, in __getitem__
    return _getitem(self, where, self.__class__._IS_MOMENTUM)  # type: ignore[arg-type]
  File "D:\extras\IRIS-HEP\.env\lib\site-packages\vector\_backends\numpy_.py", line 143, in _getitem
    return array.ObjectClass(*out.view(numpy.ndarray))  # type: ignore[misc, return-value]
TypeError: <lambda>() takes 3 positional arguments but 4 were given

Cause

Going through the stack trace I looked into the function _getitem, which was throwing an error related to the number of arguments. After adding some debug statements, I discovered that the if-else statements written to initialize azimuthal, longitudinal, and temporal are never executed when running the example.

This is due to the fact that the type of the passed array argument is AzimuthalNumpyXY (in the case of calling vec.azimuthal on a structured array having "x" and "y" as the datatypes). As array belongs to AzimuthalNumpyXY, it never has an attribute named _azimuthal_type, rather it has attributes named x and y.

The main intention here was to pass in a VectorNumpy2D, VectorNumpy3D, or VectorNumpy4D object but I can't really figure out why the __getitem method in the class GetItem is picking up AzimuthalNumpyXY.

I think the case of calling azimuthal on a VectorNumpy2D object works coincidentally, through the last else condition.

_getitem

def _getitem(
    array: typing.Union["VectorNumpy2D", "VectorNumpy3D", "VectorNumpy4D"],
    where: typing.Any,
    is_momentum: bool,
) -> typing.Union[float, FloatArray]:
    if isinstance(where, str):
        if is_momentum:
            where = _repr_momentum_to_generic.get(where, where)
        return array.view(numpy.ndarray)[where]
    else:
        out = numpy.ndarray.__getitem__(array, where)
        if not isinstance(out, numpy.void):
            return out

        azimuthal, longitudinal, temporal = None, None, None
        if hasattr(array, "_azimuthal_type"):
            azimuthal = array._azimuthal_type.ObjectClass(
                *(out[x] for x in _coordinate_class_to_names[_aztype(array)])
            )
        if hasattr(array, "_longitudinal_type"):
            longitudinal = array._longitudinal_type.ObjectClass(  # type: ignore[union-attr]
                *(out[x] for x in _coordinate_class_to_names[_ltype(array)])  # type: ignore[arg-type]
            )
        if hasattr(array, "_temporal_type"):
            temporal = array._temporal_type.ObjectClass(  # type: ignore[union-attr]
                *(out[x] for x in _coordinate_class_to_names[_ttype(array)])  # type: ignore[arg-type]
            )
        if temporal is not None:
            return array.ObjectClass(azimuthal, longitudinal, temporal)  # type: ignore[call-arg, arg-type, return-value]
        elif longitudinal is not None:
            return array.ObjectClass(azimuthal, longitudinal)  # type: ignore[call-arg, arg-type, return-value]
        elif azimuthal is not None:
            return array.ObjectClass(azimuthal)  # type: ignore[call-arg, return-value]
        else:
            return array.ObjectClass(*out.view(numpy.ndarray))  # type: ignore[misc, return-value]

A possible fix

I tried fixing this by modifying the function but ended up breaking some tests -

Modified function

def _getitem(
    array: typing.Union["VectorNumpy2D", "VectorNumpy3D", "VectorNumpy4D"],
    where: typing.Any,
    is_momentum: bool,
) -> typing.Union[float, FloatArray]:
    if isinstance(where, str):
        if is_momentum:
            where = _repr_momentum_to_generic.get(where, where)
        return array.view(numpy.ndarray)[where]
    else:
        out = numpy.ndarray.__getitem__(array, where)
        if not isinstance(out, numpy.void):
            return out
        azimuthal, longitudinal, temporal = None, None, None

        if (hasattr(array, "x") and hasattr(array, "y")):
            azimuthal = vector._backends.object_.AzimuthalObjectXY(*(out[x] for x in _coordinate_class_to_names[AzimuthalXY]))
        elif (hasattr(array, "rho") and hasattr(array, "phi")):
            azimuthal = vector._backends.object_.AzimuthalObjectRhoPhi(*(out[x] for x in _coordinate_class_to_names[AzimuthalRhoPhi]))
        
        if hasattr(array, "z"):
            longitudinal = vector._backends.object_.LongitudinalObjectZ(*(out[x] for x in _coordinate_class_to_names[LongitudinalZ]))  # type: ignore[union-attr]
        elif hasattr(array, "eta"):
            longitudinal = vector._backends.object_.LongitudinalObjectEta(*(out[x] for x in _coordinate_class_to_names[LongitudinalEta]))  # type: ignore[union-attr]
        elif hasattr(array, "theta"):
            longitudinal = vector._backends.object_.LongitudinalObjectTheta(*(out[x] for x in _coordinate_class_to_names[LongitudinalTheta]))  # type: ignore[union-attr]

        if hasattr(array, "t"):
            temporal = vector._backends.object_.TemporalObjectT(*(out[x] for x in _coordinate_class_to_names[TemporalT]))  # type: ignore[union-attr]
        elif hasattr(array, "tau"):
            temporal = vector._backends.object_.TemporalObjectTau(*(out[x] for x in _coordinate_class_to_names[TemporalTau]))  # type: ignore[union-attr]

        if temporal is not None:
            return temporal  # type: ignore[call-arg, arg-type, return-value]
        elif longitudinal is not None:
            return longitudinal  # type: ignore[call-arg, arg-type, return-value]
        elif azimuthal is not None:
            return azimuthal  # type: ignore[call-arg, return-value]
        else:
            return array.ObjectClass(*out.view(numpy.ndarray))  # type: ignore[misc, return-value]

This works somewhat fine. An example -

import vector

vec = vector.array(
    [
        (1.1, 2.1, 3.1),
        (1.2, 2.2, 3.2),
        (1.3, 2.3, 3.3),
        (1.4, 2.4, 4.4),
        (1.5, 2.5, 5.5)
    ], dtype=[("x", float), ("y", float), ("z", float)]
)
print(vec.azimuthal)

# output
# AzimuthalNumpyXY([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4),
#                 (1.5, 2.5)],
#                 dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8')])

print(vec.longitudinal)
# output
# LongitudinalNumpyZ([(3.1,), (3.2,), (3.3,), (4.4,), (5.5,)],
#                    dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8')])

Broken tests

These are most probably failing as I did a lot of hard coding inside _getitem -

=============================================================================================== FAILURES =============================================================================================== 
__________________________________________________________________________________________ test_spatial_numpy __________________________________________________________________________________________ 

    def test_spatial_numpy():
        v1 = vector._backends.numpy_.VectorNumpy3D(
            [(0.1, 0.2, 0.3)],
            dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)],
        )
        v2 = vector._backends.numpy_.VectorNumpy3D(
            [(0.4, 0.5, 0.6)],
            dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)],
        )
        out = v1.cross(v2)
        assert isinstance(out, vector._backends.numpy_.VectorNumpy3D)
        assert out.dtype.names == ("x", "y", "z")
>       assert (out[0].x, out[0].y, out[0].z) == pytest.approx((-0.03, 0.06, -0.03))
E       AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'

tests\compute\spatial\test_cross.py:55: AttributeError
__________________________________________________________________________________________ test_lorentz_numpy __________________________________________________________________________________________

    def test_lorentz_numpy():
        v1 = vector._backends.numpy_.VectorNumpy4D(
            [(0.1, 0.2, 0.3, 99)],
            dtype=[
                ("x", numpy.float64),
                ("y", numpy.float64),
                ("z", numpy.float64),
                ("t", numpy.float64),
            ],
        )
        v2 = vector._backends.numpy_.VectorNumpy4D(
            [(0.4, 0.5, 0.6, 99)],
            dtype=[
                ("x", numpy.float64),
                ("y", numpy.float64),
                ("z", numpy.float64),
                ("t", numpy.float64),
            ],
        )
        out = v1.cross(v2)
        assert isinstance(out, vector._backends.numpy_.VectorNumpy3D)
        assert out.dtype.names == ("x", "y", "z")
        assert out.tolist() == pytest.approx([(-0.03, 0.06, -0.030000000000000013)])

        for t1 in (
            "xyzt",
            "xythetat",
            "xyetat",
            "rhophizt",
            "rhophithetat",
            "rhophietat",
            "xyztau",
            "xythetatau",
            "xyetatau",
            "rhophiztau",
            "rhophithetatau",
            "rhophietatau",
        ):
            for t2 in (
                "xyzt",
                "xythetat",
                "xyetat",
                "rhophizt",
                "rhophithetat",
                "rhophietat",
                "xyztau",
                "xythetatau",
                "xyetatau",
                "rhophiztau",
                "rhophithetatau",
                "rhophietatau",
            ):
                transformed1, transformed2 = (
                    getattr(v1, "to_" + t1)(),
                    getattr(v2, "to_" + t2)(),
                )
                out = transformed1.cross(transformed2)
                assert isinstance(out, vector._backends.numpy_.VectorNumpy3D)
                assert out.dtype.names == ("x", "y", "z")
>               assert (out[0].x, out[0].y, out[0].z) == pytest.approx((-0.03, 0.06, -0.03))
E               AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'

tests\compute\spatial\test_cross.py:186: AttributeError
__________________________________________________________________________________________ test_spatial_numpy __________________________________________________________________________________________ 

    def test_spatial_numpy():
        vec = vector._backends.numpy_.VectorNumpy3D(
            [(0.1, 0.2, 0.3)],
            dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)],
        )
        out = vec.rotateX(0.25)
        assert isinstance(out.azimuthal, vector._methods.AzimuthalXY)
        assert isinstance(out.longitudinal, vector._methods.LongitudinalZ)
>       assert out[0].x == pytest.approx(0.1)
E       AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'

tests\compute\spatial\test_rotateX.py:43: AttributeError
__________________________________________________________________________________________ test_lorentz_numpy __________________________________________________________________________________________

    def test_lorentz_numpy():
        vec = vector._backends.numpy_.VectorNumpy4D(
            [(0.1, 0.2, 0.3, 99)],
            dtype=[
                ("x", numpy.float64),
                ("y", numpy.float64),
                ("z", numpy.float64),
                ("t", numpy.float64),
            ],
        )
        out = vec.rotateX(0.25)
        assert isinstance(out.azimuthal, vector._methods.AzimuthalXY)
        assert isinstance(out.longitudinal, vector._methods.LongitudinalZ)
>       assert out[0].x == pytest.approx(0.1)
E       AttributeError: 'TemporalObjectT' object has no attribute 'x'

tests\compute\spatial\test_rotateX.py:106: AttributeError
__________________________________________________________________________________________ test_spatial_numpy __________________________________________________________________________________________ 

    def test_spatial_numpy():
        vec = vector._backends.numpy_.VectorNumpy3D(
            [(0.1, 0.2, 0.3)],
            dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)],
        )
        out = vec.rotateY(0.25)
        assert isinstance(out.azimuthal, vector._methods.AzimuthalXY)
        assert isinstance(out.longitudinal, vector._methods.LongitudinalZ)
>       assert out[0].x == pytest.approx(0.17111242994742137)
E       AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'

tests\compute\spatial\test_rotateY.py:43: AttributeError
__________________________________________________________________________________________ test_lorentz_numpy __________________________________________________________________________________________

    def test_lorentz_numpy():
        vec = vector._backends.numpy_.VectorNumpy4D(
            [(0.1, 0.2, 0.3, 99)],
            dtype=[
                ("x", numpy.float64),
                ("y", numpy.float64),
                ("z", numpy.float64),
                ("t", numpy.float64),
            ],
        )
        out = vec.rotateY(0.25)
        assert isinstance(out.azimuthal, vector._methods.AzimuthalXY)
        assert isinstance(out.longitudinal, vector._methods.LongitudinalZ)
>       assert out[0].x == pytest.approx(0.17111242994742137)
E       AttributeError: 'TemporalObjectT' object has no attribute 'x'

tests\compute\spatial\test_rotateY.py:106: AttributeError
__________________________________________________________________________________________ test_spatial_numpy __________________________________________________________________________________________ 

    def test_spatial_numpy():
        axis = vector._backends.numpy_.VectorNumpy3D(
            [(0.1, 0.2, 0.3)],
            dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)],
        )
        vec = vector._backends.numpy_.VectorNumpy3D(
            [(0.4, 0.5, 0.6)],
            dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)],
        )
        out = vec.rotate_axis(axis, 0.25)
        assert isinstance(out, vector._backends.numpy_.VectorNumpy3D)
        assert out.dtype.names == ("x", "y", "z")
>       assert out[0].x == pytest.approx(0.37483425404335763)
E       AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'

tests\compute\spatial\test_rotate_axis.py:59: AttributeError
__________________________________________________________________________________________ test_lorentz_numpy __________________________________________________________________________________________ 

    def test_lorentz_numpy():
        axis = vector._backends.numpy_.VectorNumpy4D(
            [(0.1, 0.2, 0.3, 99)],
            dtype=[
                ("x", numpy.float64),
                ("y", numpy.float64),
                ("z", numpy.float64),
                ("t", numpy.float64),
            ],
        )
        vec = vector._backends.numpy_.VectorNumpy4D(
            [(0.4, 0.5, 0.6, 99)],
            dtype=[
                ("x", numpy.float64),
                ("y", numpy.float64),
                ("z", numpy.float64),
                ("t", numpy.float64),
            ],
        )
        out = vec.rotate_axis(axis, 0.25)
        assert isinstance(out, vector._backends.numpy_.VectorNumpy4D)
        assert out.dtype.names == ("x", "y", "z", "t")
>       assert out[0].x == pytest.approx(0.37483425404335763)
E       AttributeError: 'TemporalObjectT' object has no attribute 'x'

tests\compute\spatial\test_rotate_axis.py:163: AttributeError
=========================================================================================== warnings summary =========================================================================================== 
c:\users\saransh\saransh_softwares\python_3.9\lib\site-packages\pyreadline\py3k_compat.py:8
  c:\users\saransh\saransh_softwares\python_3.9\lib\site-packages\pyreadline\py3k_compat.py:8: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3, and in 3.10 it will stop working
    return isinstance(x, collections.Callable)

-- Docs: https://docs.pytest.org/en/stable/warnings.html
======================================================================================= short test summary info ======================================================================================== 
SKIPPED [2404] tests\test_compute_features.py:99: Unsupported Python version 3.9 (canonic 3.9.0beta5)
FAILED tests/compute/spatial/test_cross.py::test_spatial_numpy - AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'
FAILED tests/compute/spatial/test_cross.py::test_lorentz_numpy - AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'
FAILED tests/compute/spatial/test_rotateX.py::test_spatial_numpy - AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'
FAILED tests/compute/spatial/test_rotateX.py::test_lorentz_numpy - AttributeError: 'TemporalObjectT' object has no attribute 'x'
FAILED tests/compute/spatial/test_rotateY.py::test_spatial_numpy - AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'
FAILED tests/compute/spatial/test_rotateY.py::test_lorentz_numpy - AttributeError: 'TemporalObjectT' object has no attribute 'x'
FAILED tests/compute/spatial/test_rotate_axis.py::test_spatial_numpy - AttributeError: 'LongitudinalObjectZ' object has no attribute 'x'
FAILED tests/compute/spatial/test_rotate_axis.py::test_lorentz_numpy - AttributeError: 'TemporalObjectT' object has no attribute 'x'