How to create a new object with awkward?

[ZhenxuanZhang-Jensen]

Hi experts,
I read the jets kinematic from a root file, and save it as an awkward array. And it's a jagged array, every event has different numbers of jets. How can I combine two of those jets in each event to be a new object(W boson), and then choose the one closest to the W inv mass in each event and save it in the event array as a new object called "W"?

cat_4jets_jets_pt = events['SelectedJet_pt'][ak.num(events['SelectedJet_pt']) >= 4]
cat_4jets_jets_eta = events['SelectedJet_eta'][ak.num(events['SelectedJet_eta']) >= 4]
cat_4jets_jets_phi = events['SelectedJet_phi'][ak.num(events['SelectedJet_phi']) >= 4]
cat_4jets_jets_mass = events['SelectedJet_mass'][ak.num(events['SelectedJet_mass']) >= 4]

cat_4jets_jets_LV =  vector.zip({
    "pt": cat_4jets_jets_pt,
    "phi": cat_4jets_jets_eta,
    "eta": cat_4jets_jets_phi,
    "M": cat_4jets_jets_mass,
})

[agoose77]

First things's first, I'd ak.zip the fields of the jet (pt, eta, ...) into their own record:

import awkward as ak

jet = ak.zip({
    'pt': events.SelectedJet_pt,
    'eta': events.SelectedJet_eta,
    'phi': events.SelectedJet_phi,
    'mass': events.SelectedJet_mass,
})

The vector library gives useful helper classes that add special behaviours to Awkward Arrays. We can use the Momentum4D class to get a four-vector behaviour:

import vector

jet = vector.zip({
    'pt': events.SelectedJet_pt,
    'eta': events.SelectedJet_eta,
    'phi': events.SelectedJet_phi,
    'mass': events.SelectedJet_mass,
})

Note that, vector.zip is just shorthand for ak.zip + some logic that guesses which behaviour you want. You can also use

vector.register_awkward()

jet = ak.zip({
    'pt': events.SelectedJet_pt,
    'eta': events.SelectedJet_eta,
    'phi': events.SelectedJet_phi,
    'mass': events.SelectedJet_mass,
}, with_name="Momentum4D")

To solve your problem, we don't need this behaviour, because you just want the invariant mass. But, it's helpful to use it in case you later decide to do something more complex. The subsequent solutions will make use of Momentum4D in order to add jets (four vectors) together.

four_jet = jet[ak.num(jet, axis=-1) >= 4]

To explore all of the pairwise combinations of the jets, you can useak.combinations:

pair = ak.combinations(four_jet, axis=1,  n=2)

Then you can add together the four vectors to compute the boson vector

boson = pair.slot0 + pair.slot1

With the boson vector we can compute the absolute mass difference between each boson candidate and the known W boson mass:

w_boson_mass = ...
delta_mass = np.abs(boson.mass - w_boson_mass)

Then, finally, we can pull out the boson with the shortest mass.

i_best = ak.argmin(
    delta_mass, 
    axis=-1, 
    keepdims=True
)
w_boson = boson[i_best]

[ZhenxuanZhang-Jensen]

First of all, thanks for the answers and the clues, it helps me to get this straight in my mind now. But in this way, indeed we can have a W boson, but as you know, we get 4 jets as an array to finally form the Higgs boson candidate, so after finding the closest to 80GeV W boson, we need to find two jets out of the remained jets to form a Higgs boson. I'm thinking maybe change the ak.combinations numbers to 4 or so? Do you have any ideas? thanks.

[jpivarski]

ak.combinations(jets, 4) is what you want, though I hope you have enough memory for all the combinations. There was another question recently in which ak.combinations(jets, 7) was needed, and for that I suggested Numba, since it might be hard to keep all of those combinations in memory at once.

Aha—that was you! #1500 (reply in thread) Sorry for not recognizing at first.

[ZhenxuanZhang-Jensen]

Yeah, that is also the reason I try to use vector in Numba and also face some issues on it. About the ak.combinations(jets,4), I still get some confusion. Can we choose the two jets of the four in the combination which is closest to the W boson mass just like using ak.combination(jets,2) and slot0,slot1 then form choose the rest of them closest to higgs invM? I mean first, we need to choose the closest to 80 Gev two jets and then we need to choose two jets in the rest of all jets to form as the closest to 125 GeV Higgs candidate.

[jpivarski]

Just linking everything together, so that we can see what you've tried and what's working/not working:

• how to add two LorentzVector? vector#205
• How to use vector in a numba function? vector#206

[ZhenxuanZhang-Jensen]

yes, thanks Jim, it basically the problems I met when I do the same analysis with uproot,awkward,vector and numba