Invariant mass calculations with the vector package

[JanFSchulte]

Hi,

to preserve an email exchange I just had with @jpivarski:

Question:

I am currently trying to update the LPC Muon HATS for this year. The exercise is based on nanoAOD and was based on awkward 0.x and used the uproot_methods package to do TLorentzVector-related stuff. To not teach updated things, I am trying to get it work with awkward 1.x and using the vector package as a replacement for the TLorentzVectors.

Last years version of the exercise can be found here: https://github.com/Taiantuia/MuonHATSatLPC2020/blob/master/Exercise-2-Muon-momentum-scale-and-resolution-corrections.ipynb

I am basically trying to replace these two lines:

two_muon_p4=uproot_methods.TLorentzVectorArray.from_ptetaphim(two_opposite_good_muons.muon_pt, two_opposite_good_muons.muon_eta, two_opposite_good_muons.muon_phi, two_opposite_good_muons.muon_mass)

dimuon_p4=two_muon_p4[:,0]+two_muon_p4[:,1

Based on the vector documentation, I am creating an awkward array using

two_muon_p4=vector.awk({"pt":two_opposite_good_muons["Muon_pt"], "eta":two_opposite_good_muons["Muon_eta"],
"phi":two_opposite_good_muons["Muon_phi"], "M":two_opposite_good_muons["Muon_mass"]})

And from what I can see, I get the expected structure of nested arrays where there are two vectors in each sub-array. Now I am trying to combine them as before with

dimuon_p4=two_muon_p4[:,0]+two_muon_p4[:,1]

which fails with

ValueError: cannot broadcast records in this type of operation

I am currently stuck at this point as it is not clear to me how the correct syntax would be doing this with vector objects. Could you point me in the right direction?

Thanks a lot!
Jan

And answer:

An answer that could be moved there is this: we're still working out what would be the best ways to access it, which is something that has to be done with users, to see what works well in their workflows. I really don't like the name "vector.awk", and there are ways of using that function that builds incorrect Awkward Array structures. (To keep using this function, we'd at least have to have it sanitize the input so that what it builds is either right or an error.)

What about the following technique? Vector's Awkward backend is just a "behavior" in Awkward Array (see documentation), which can be installed globally for all Awkward Arrays with "vector.register_awkward()".

import awkward as ak
>>> import vector
>>> vector.register_awkward()

Then, any array with the right-named record and right-named fields would have all of the Lorentz vector functions. This array is pretty close to what you already have:

muons = ak.Array([
... [{"pt": 5.0, "eta": 1.5, "phi": 0.1, "mass": 0.105},
... {"pt": 10.0, "eta": -1.5, "phi": 0.2, "mass": 0.105}],
... [{"pt": 5.0, "eta": 1.5, "phi": 0.1, "mass": 0.105},
... {"pt": 10.0, "eta": -1.5, "phi": 0.2, "mass": 0.105}]])
>>> muons
<Array [[{pt: 5, eta: 1.5, ... mass: 0.105}]] type='2 * var * {"pt": float64, "e...'>
>>> muons.type
2 * var * {"pt": float64, "eta": float64, "phi": float64, "mass": float64}

If you need to rename fields, there's ak.zip. It's a generic Awkward Array because the records don't have names. The Lorentz vector methods have all been registered to the name "Momentum4D", so we could just rename it that. (with_name is an argument of ak.zip, too.)

muon_vectors = ak.Array(muons, with_name="Momentum4D")
>>> muon_vectors
<MomentumArray4D [[{pt: 5, eta: 1.5, ... mass: 0.105}]] type='2 * var * Momentum...'>
>>> muon_vectors.type
2 * var * Momentum4D["pt": float64, "eta": float64, "phi": float64, "mass": float64]

Now it knows about all of the Lorentz vector methods:

>>> muon_vectors[:, 0] + muon_vectors[:, 1]
<MomentumArray4D [{rho: 15, phi: 0.167, ... tau: 30.1}] type='2 * Momentum4D["rh...'>
>>> (muon_vectors[:, 0] + muon_vectors[:, 1]).pt
<Array [15, 15] type='2 * float64'>
>>> muon_vectors[:, 0].deltaR(muon_vectors[:, 1])
<Array [3, 3] type='2 * float64'>

etc. By contrast, the original muons raises exceptions if you try to do these things (can't broadcast "+" through records, doesn't have a method or field named "deltaR"). As long as the record name is "Momentum4D," the field names are acceptable (x, y, rho, z, t, tau, etc. are geometrical names; px, py, pt, pz, energy, mass, etc. are momentum synonyms—interchangeable on an array that has been declared "Momentum4D", rather than "Vector4D"), and as long as the behaviors have been registered, these records are Lorentz vectors.

[JanFSchulte]

Happy to report that this worked for me!

[jpivarski]

Great! Maybe the tutorials should be based on this method. The trouble is that everyone who has been involved in the Vector project so far is busy with other commitments. If there's anybody out there who uses Lorentz vectors every day and wants to clean the rough edges off this project, including its front-page README and other documentation, it would be highly appreciated!