ak.sum over records?

[rkansal47]

Hi, is it possible to sum over records using ak.sum() while respecting the records' 'add' behaviour? Specifically I'm using the PtEtaPhiMLorentzVector record from Coffea and I am seeing that performing an ak.sum() over an array of of Vectors is not equivalent to summing them manually (i.e. with the '+' operation).

e.g. with such an array:
test_vecs = ak.zip({ "pt": [1, 2], "eta": [0.1, 0.2], "phi": [0.1, 0.2], "mass": [0, 0] }, with_name="PtEtaPhiMLorentzVector")

ak.sum(test_vecs, axis=0) just sums each individual fields' arrays whereas I'm hoping to get the same output as doing test_vecs[0] + test_vecs[1] (which would add them using the defined behaviour).

[jpivarski]

Currently, it is not, though I can see how that would be nice. The reducers (like sum) are handled deep in C++ to make axis != -1 possible, but behaviors like PtEtaPhiMLorentzVector are only known to the Python layer.

In fact, it might be necessary to prevent ak.sum from descending into records because it just adds components like this, and that would be undesirable (and surprising) for non-Cartesian coordinates like this.

It may be possible to get what you want—sums of variable numbers of Lorentz vectors without having to pick out each one—but it's not going to be an easy fix. @nsmith- asked about this a while ago.

[nsmith-]

As a workaround for lack of ak.sum support, all the coffea vector objects have a .sum() method. See https://coffeateam.github.io/coffea/api/coffea.nanoevents.methods.vector.LorentzVector.html#coffea.nanoevents.methods.vector.LorentzVector.sum
In your example, however, test_vecs.sum() fails due to #758. For a jagged test_vecs:

>>> test2 = ak.unflatten(test_vecs, [2])
>>> test2.sum()[0]
<LorentzVectorRecord ... y: 0.497, z: 0.503, t: 3.05} type='LorentzVector["x": f...'>
>>> test2.sum()[0].x
2.955137320960509