Awkward array of Lorentz Vectors from Awkward arrays of components

[MoAly98]

Hi,

I have the following awkward array:

akarr = {"px": [1,2,3], "py": [4,5,6], "pz": [7,8,9], "E": [10,11,12]}

And I would like to add to this array another field which holds the four momenta from these values in akarr. Does awkward support any of the LorentzVector object from other sci-hep packages? if so, how would I add this lorentz_vec field?

Thank you very much.
Mo

[agoose77]

The vector package provides a set of LorentzXXX classes, which exposes a suite of methods and properties relating specifically to Lorentz vectors.

vector can take existing arrays, and figure out which kind of array to return to you according to the fields that it has. For example, given your akarr:

>>> import awkward as ak
>>> import vector
>>> akarr = vector.zip({"px": [1,2,3], "py": [4,5,6], "pz": [7,8,9], "E": [10,11,12]})
>>> akarr
<MomentumArray4D [{x: 1, y: 4, z: 7, ... z: 9, t: 12}] type='3 * Momentum4D["x":...'>

Note that the array has more than just the Lorentz methods: it is a MomentumArray4D which pulls in lots of other classes:

>>> akarr.__class__.__mro__
(vector._backends.awkward_.MomentumArray4D,
 vector._backends.awkward_.MomentumAwkward4D,
 vector._methods.LorentzMomentum,
 vector._methods.SpatialMomentum,
 vector._methods.PlanarMomentum,
 vector._methods.Momentum,
 vector._methods.MomentumProtocolLorentz,
 vector._backends.awkward_.VectorAwkward4D,
 vector._backends.awkward_.VectorAwkward,
 vector._methods.Lorentz,
 vector._methods.Spatial,
 vector._methods.Planar,
 vector._methods.Vector4D,
 vector._methods.Vector,
 vector._methods.VectorProtocolLorentz,
 vector._methods.MomentumProtocolSpatial,
 vector._methods.VectorProtocolSpatial,
 vector._methods.MomentumProtocolPlanar,
 vector._methods.VectorProtocolPlanar,
 vector._methods.VectorProtocol,
 awkward.highlevel.Array,
 numpy.lib.mixins.NDArrayOperatorsMixin,
 collections.abc.Iterable,
 collections.abc.Sized,
 object)

The vector.zip function is just a convenience function. You can do this by hand by first registering the Vector behaviours, so that Awkward can find them:

>>> vector.register_awkward()

Then, using the with_name field in ak.zip to inform Awkward that the array should have the Momentum4D behavior:

>>> ak.zip({"px": [1,2,3], "py": [4,5,6], "pz": [7,8,9], "E": [10,11,12]}, with_name="Momentum4D")
<MomentumArray4D [{px: 1, py: 4, pz: 7, ... E: 12}] type='3 * Momentum4D["px": i...'>

If you already have an array, and you just want to make it behave like Momentum4D, then you can use vector.awk:

>>> arr = ak.zip({"px": [1,2,3], "py": [4,5,6], "pz": [7,8,9], "E": [10,11,12]})
>>> vec = vector.awk(arr)
>>> vec
<MomentumArray4D [{x: 1, y: 4, z: 7, ... z: 9, t: 12}] type='3 * Momentum4D["x":...'>

Like with vector.zip, vector.awk just figures out the name Momentum4D to use (from the fields in your array). You can also do this by hand:

>>> arr = ak.zip({"px": [1,2,3], "py": [4,5,6], "pz": [7,8,9], "E": [10,11,12]})
>>> vec = ak.with_name(arr, "Momentum4D")
>>> vec

I'm not entirely clear on what field you're trying to add (your vector is already a four-vector), but I imagine that you're most interested in using Vector, so that's the question I've answered :)