Store additional information in TLorentzVector

[dguest]

I asked earlier about constructing the Lorentz vector objects directly from numpy arrays. That works really well! But I'd also like to be able to add additional information to the objects, i.e. in some cases we have other jet characteristics like b-tagging scores. Is there some method to append additional information (which is stored in another array) to each object?

[jpivarski]

Yes. A TLorentzVectorArray contains a Table, so you should be able to attach or remove any other arrays to it, as long as they fit (same number of elements, and if jagged, same number in each subarray).

a = TLorentzVectorArray.from_ptetaphim(pt, eta, phi, mass)
a["isolation"] = isolation

and so on.

[dguest]

OK, that's ridiculously easy.

[dguest]

OK, while this obviously isn't the recommended way to do things, I noticed that when I looped over the jets the iterator seemed to return a regular TLorentzVector without the isolation attached. Is there some way to get both together? I noticed that the following seems to work:

for vec, iso in zip(vecarray, vecarray.content['isolation']):
    print(vec, iso)

but it's a bit clunky.

[jpivarski]

To get that to work, I should have made the TLorentzVector object (singular) a mixin of Table.Row, instead of a separate class. At some point, I should do that so that you can get isolation from vec["isolation"].

I'll reopen this to remind myself to go back and do that.

[nsmith-]

Another approach is to nest the ObjectArray inside a Table, e.g.

import awkward as ak
jets = ak.Table(
    p4=TLorentzVectorArray.from_ptetaphim(pt, eta, phi, mass),
    iso=isolation_vec,
)

Then it should carry around the fields, with the small inconvenience of having to access the lorentz vector properties through the name: jets.p4.pt, etc.