Summing over array dimensions does not work properly

[masonproffitt]

Summing doesn't seem supported at all with numpy arrays:

>>> np.sum(vector.array({"x": [1], "y": [1]}))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 5, in sum
TypeError: no implementation found for 'numpy.sum' on types that implement __array_function__: [<class 'vector._backends.numpy_.VectorNumpy2D'>]
>>> vector.array({"x": [1], "y": [1]}).sum()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/miniconda3/envs/iris-hep/lib/python3.9/site-packages/numpy/core/_methods.py", line 47, in _sum
    return umr_sum(a, axis, dtype, out, keepdims, initial, where)
TypeError: operand type(s) all returned NotImplemented from __array_ufunc__(<ufunc 'add'>, 'reduce', VectorNumpy2D([(1., 1.)], dtype=[('x', '<f8'), ('y', '<f8')]), axis=None, dtype=None, keepdims=False, where=True): 'VectorNumpy2D'

With ak.Array, you can get results with ak.sum(), but they don't make sense for vectors:

>>> ak.sum(vector.Array({"x": [1], "y": [1]}))
2

Using the axis parameter works fine for Cartesian coordinates:

>>> ak.sum(vector.Array({"x": [1, 2], "y": [3, 4]}), axis=0)
<Record {x: 3, y: 7} type='{"x": int64, "y": int64}'>

but not for any other cases:

>>> ak.sum(vector.Array({"pt": [1, 1], "phi": [0, np.pi]}), axis=0)
<Record {rho: 2, phi: 3.14} type='{"rho": int64, "phi": float64}'>

[jpivarski]

@lukasheinrich asked about this for Awkward Arrays of Vectors, rather than NumPy arrays of Vectors. (I haven't been able to find the reference.) To support np.sum/ak.sum for Awkward Arrays of Vectors, we'll need to add a new mechanism for adding behavior to Awkward Array—currently, you can't overload reducers, but that would be a nice feature. To support np.sum for NumPy arrays of Vectors, we'd just have to implement __array_function__. That's doable.

[lukasheinrich]

just linking the issue (which was transferred from vector to awkward) scikit-hep/awkward#824

[Saransh-cpp]

I was going through this issue and I tried implementing a rough version of np.sum() for 2D Numpy vectors with some basic arguments. It works!

If we decide to implement this, would it be required to add all the original arguments of np.sum()? Or only the ones which seem useful from the point of view of vector?

Implementation for a single Numpy vector class -

HANDLED_FUNCTIONS_FOR_NUMPY_2D = {}

def implements_numpy_2d(numpy_function):
    """Register an __array_function__ implementation for VectorNumpy2D objects."""
    def decorator(func):
        HANDLED_FUNCTIONS_FOR_NUMPY_2D[numpy_function] = func
        return func
    return decorator

@implements_numpy_2d(numpy.sum)
def sum_vector_numpy_2d(vec, **kwargs):
    axis = kwargs["axis"] if "axis" in kwargs else None
    dtype = kwargs["dtype"] if "dtype" in kwargs else None
    initial = kwargs["initial"] if "initial" in kwargs else 0.

    if axis == None:
        return initial + numpy.sum(vec.x, dtype=dtype) + numpy.sum(vec.y, dtype=dtype)
    elif axis == 0:
        return initial + numpy.array([numpy.sum(vec.x), numpy.sum(vec.y)], dtype=dtype)
    elif axis == 1:
        return initial + vec.x + vec.y


class VectorNumpy2D(VectorNumpy, Planar, Vector2D, FloatArray):  # type: ignore[misc]
    ObjectClass = vector._backends.object_.VectorObject2D
    _IS_MOMENTUM = False

    _azimuthal_type: typing.Union[
        typing.Type[AzimuthalNumpyXY], typing.Type[AzimuthalNumpyRhoPhi]
    ]

    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS_FOR_NUMPY_2D:
            return NotImplemented
        if not all(issubclass(t, VectorNumpy2D) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS_FOR_NUMPY_2D[func](*args, **kwargs)

I also tried generalizing the function for all Numpy Vectors (excluding the Momentum ones for now), and it works too! The rough implementation looks a bit ugly, so would it be better to keep separate functions for every Numpy vector?

A general implementation for non-momentum Numpy vector classes -

HANDLED_FUNCTIONS_FOR_NUMPY = {}

def implements_numpy(numpy_function):
    """Register an __array_function__ implementation for VectorNumpy2D objects."""
    def decorator(func):
        HANDLED_FUNCTIONS_FOR_NUMPY[numpy_function] = func
        return func
    return decorator

@implements_numpy(numpy.sum)
def sum_vector_numpy(vec, **kwargs):
    axis = kwargs["axis"] if "axis" in kwargs else None
    dtype = kwargs["dtype"] if "dtype" in kwargs else None
    initial = kwargs["initial"] if "initial" in kwargs else 0.

    if axis == None:
        if vec._azimuthal_type == AzimuthalNumpyXY:
            ret_sum = numpy.sum(vec.x, dtype=dtype) + numpy.sum(vec.y, dtype=dtype)
        else:
            ret_sum = numpy.sum(vec.rho, dtype=dtype) + numpy.sum(vec.phi, dtype=dtype)

        if isinstance(vec, (Vector3D, Vector4D)):
            if vec._longitudinal_type == LongitudinalNumpyZ:
                ret_sum += numpy.sum(vec.z, dtype=dtype)
            elif vec._longitudinal_type == LongitudinalNumpyEta:
                ret_sum += numpy.sum(vec.eta, dtype=dtype)
            else:
                ret_sum += numpy.sum(vec.theta, dtype=dtype)

        if isinstance(vec, Vector4D):
            if vec._temporal_type == TemporalNumpyT:
                ret_sum += numpy.sum(vec.t, dtype=dtype)
            else:
                ret_sum += numpy.sum(vec.tau, dtype=dtype)

        return ret_sum + initial
    
    elif axis == 0:
        ret_sum = []
        if vec._azimuthal_type == AzimuthalNumpyXY:
            ret_sum.extend([numpy.sum(vec.x, dtype=dtype), numpy.sum(vec.y, dtype=dtype)])
        else:
            ret_sum.extend([numpy.sum(vec.rho, dtype=dtype), numpy.sum(vec.phi, dtype=dtype)])
        
        if isinstance(vec, (Vector3D, Vector4D)):
            if vec._longitudinal_type == LongitudinalNumpyZ:
                ret_sum.append(numpy.sum(vec.z, dtype=dtype))
            elif vec._longitudinal_type == LongitudinalNumpyEta:
                ret_sum.append(numpy.sum(vec.eta, dtype=dtype))
            else:
                ret_sum.append(numpy.sum(vec.theta, dtype=dtype))

        if isinstance(vec, Vector4D):
            if vec._temporal_type == TemporalNumpyT:
                ret_sum.append(numpy.sum(vec.t, dtype=dtype))
            else:
                ret_sum.append(numpy.sum(vec.tau, dtype=dtype))

        return numpy.array(ret_sum) + initial

    elif axis == 1:
        if vec._azimuthal_type == AzimuthalNumpyXY:
            ret_sum = vec.x + vec.y
        else:
            ret_sum = vec.rho + vec.phi

        if isinstance(vec, (Vector3D, Vector4D)):
            if vec._longitudinal_type == LongitudinalNumpyZ:
                ret_sum += vec.z
            elif vec._longitudinal_type == LongitudinalNumpyEta:
                ret_sum += vec.eta
            else:
                ret_sum += vec.theta

        if isinstance(vec, Vector4D):
            if vec._temporal_type == TemporalNumpyT:
                ret_sum += vec.t
            else:
                ret_sum += vec.tau

        return ret_sum + initial


class VectorNumpy2D(VectorNumpy, Planar, Vector2D, FloatArray):  # type: ignore[misc]
    ObjectClass = vector._backends.object_.VectorObject2D
    _IS_MOMENTUM = False

    _azimuthal_type: typing.Union[
        typing.Type[AzimuthalNumpyXY], typing.Type[AzimuthalNumpyRhoPhi]
    ]

    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS_FOR_NUMPY:
            return NotImplemented
        if not all(issubclass(t, VectorNumpy2D) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS_FOR_NUMPY[func](*args, **kwargs)
.
.
.

class VectorNumpy3D(VectorNumpy, Spatial, Vector3D, FloatArray):  # type: ignore[misc]
    ObjectClass = vector._backends.object_.VectorObject3D
    _IS_MOMENTUM = False

    _azimuthal_type: typing.Union[
        typing.Type[AzimuthalNumpyXY], typing.Type[AzimuthalNumpyRhoPhi]
    ]
    _longitudinal_type: typing.Union[
        typing.Type[LongitudinalNumpyZ],
        typing.Type[LongitudinalNumpyTheta],
        typing.Type[LongitudinalNumpyEta],
    ]

    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS_FOR_NUMPY:
            return NotImplemented
        if not all(issubclass(t, VectorNumpy3D) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS_FOR_NUMPY[func](*args, **kwargs)
.
.
.

class VectorNumpy4D(VectorNumpy, Lorentz, Vector4D, FloatArray):  # type: ignore[misc]
    ObjectClass = vector._backends.object_.VectorObject4D
    _IS_MOMENTUM = False

    _azimuthal_type: typing.Union[
        typing.Type[AzimuthalNumpyXY], typing.Type[AzimuthalNumpyRhoPhi]
    ]
    _longitudinal_type: typing.Union[
        typing.Type[LongitudinalNumpyZ],
        typing.Type[LongitudinalNumpyTheta],
        typing.Type[LongitudinalNumpyEta],
    ]
    _temporal_type: typing.Union[
        typing.Type[TemporalNumpyT],
        typing.Type[TemporalNumpyTau],
    ]

    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS_FOR_NUMPY:
            return NotImplemented
        if not all(issubclass(t, VectorNumpy4D) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS_FOR_NUMPY[func](*args, **kwargs)

[jpivarski]

That's great! The main one to get is axis, and it looks like you got that.

It doesn't address the problem for Awkward Array, in which the issue is telling ak.Arrays to use a custom implementation, not just writing the implementation. That will require changes in Awkward itself, and ultimately it will only be possible for Awkward version 2.

But having this addresses the problem for NumPy arrays of vectors, and therefore it's a step forward.