How to obtain matched genJets into an matching array with Jet

[green-cabbage]

Hello,

I am trying to double check that delta_r of Jet and GenJet saved in Jet.genJetIdx is indeed less than 0.4. However, I am having a hard time writing this in code. What I would like to do is something like:

import awkward as ak
from coffea.nanoevents import NanoEventsFactory, NanoAODSchema

events = NanoEventsFactory.from_root(.....)
jets = events.Jet
genjets = events.GenJet
has_matched_gj = jets.genJetIdx >= 0
matched_jets = jets[has_matched_gj]
matched_gjets = genjets.eta[matched_jets.genJetIdx]
deltaR = matched_jets.delta_r(matched_gjets)
print(ak.all(deltaR < 0.4))

However, I get errors in the indexing of the genjets something to the effect of:

IndexError: cannot slice ListArray (of length 2000) with [[1], [1, 2], [2], [1], ..., [1, 2, 4, 3, 5], [], [2, 1, 3, 4], [1, 3, 2]]: index out of range while attempting to get index 4

I am 95% certain that I am missing some neat awkward array trick, but I couldn't find them in the documentation. Some hints would be much appreciated!

Thank you in advance,

Hyeon-Seo

[ikrommyd]

I believe your problem is in the matched_gjets = genjets.eta[matched_jets.genJetIdx] line. matched_jets.genJetIdx has the same dimensionality as events.Jet so you can't mask events.GenJet with it, only events.Jet.

What I would personally use is something that is alreayd built in coffea. And that is events.Jet.matched_gen. This will give you an awkward array of GenJets with the same dimensionality as events.Jet where Jets that are associated with a GenJet through the genJetIdx index are replaced by the corresponding GenJet and None otherwise. That way you can skip doing all that yourself.
Then you could calculate the deltaR between the matched genjets and jet like events.Jet.delta_r(events.Jet.matched_gen).

In [6]: events = NanoEventsFactory.from_root({"../../coffea_dev/coffea/tests/samples/nano_dy.root": "Events"}, delayed=False).events()
/Users/iason/miniforge3/envs/egamma_dev/lib/python3.10/site-packages/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for LowPtElectron_electronIdx => Electron
  warnings.warn(
/Users/iason/miniforge3/envs/egamma_dev/lib/python3.10/site-packages/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for LowPtElectron_genPartIdx => GenPart
  warnings.warn(
/Users/iason/miniforge3/envs/egamma_dev/lib/python3.10/site-packages/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for LowPtElectron_photonIdx => Photon
  warnings.warn(
/Users/iason/miniforge3/envs/egamma_dev/lib/python3.10/site-packages/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for FatJet_genJetAK8Idx => GenJetAK8
  warnings.warn(

In [7]: events.Jet
Out[7]: <JetArray [[Jet, ..., Jet], ...] type='40 * var * Jet[area: float32[paramet...'>

In [8]: events.Jet.matched_gen
Out[8]: <PtEtaPhiMCollectionArray [[{eta: 1.55, ...}, ..., None], ...] type='40 * v...'>

In [9]: events.GenJet
Out[9]: <PtEtaPhiMCollectionArray [[{eta: 1.55, ...}, ..., {...}], ...] type='40 * ...'>

In [10]: events.Jet.delta_r(events.Jet.matched_gen)
Out[10]: <Array [[0.011, 0.085, 0.152, None, None], ...] type='40 * var * ?float32'>

And with this events.Jet.delta_r(events.Jet.matched_gen) array you can check if it's smaller than 0.4 where it's not None as None means the Jet isn't matched by index to a GenJet like

In [11]: ak.all(events.Jet.delta_r(events.Jet.matched_gen) < 0.4)
Out[11]: True

[green-cabbage]

Thanks @ikrommyd for the succinct reply! From your answer, I looked at the definition of the matched_gen and therein lies the code that I was looking for:

@dask_property
def matched_gen(self):
    return self._events().GenJet._apply_global_index(self.genJetIdxG)

Therefore, if I use the above definition, I get the same result as your method:

import dask_awkward as dak
import awkward as ak
from coffea.nanoevents import NanoEventsFactory, NanoAODSchema
from distributed import LocalCluster, Client

file_dict = {"root://eos.cms.rcac.purdue.edu//store/mc/RunIIAutumn18NanoAODv6/DYJetsToLL_M-105To160_TuneCP5_PSweights_13TeV-amcatnloFXFX-pythia8/NANOAODSIM/Nano25Oct2019_102X_upgrade2018_realistic_v20-v1/100000/1E5C1F4B-FE04-624F-82FD-C4273792B1D1.root": {"object_path": "Events", "steps": [[0, 1000], [1000, 2000]]}}


client = Client(n_workers=1,  threads_per_worker=1, processes=True)
events = NanoEventsFactory.from_root(
    file_dict,
    schemaclass=NanoAODSchema,
).events()

def checkGenIdxRecommended(events):
    delta_rs = events.Jet.delta_r(events.Jet.matched_gen) 
    delta_rs =  delta_rs.compute()
    # print(delta_r_test)
    print(f"checkGenIdxRecommended: {ak.all(delta_rs <= 0.4)}")


def checkGenIdxManual(events):
    delta_rs = events.Jet.delta_r(events.GenJet._apply_global_index(events.Jet.genJetIdxG)) 
    delta_rs =  delta_rs.compute()
    # print(delta_r_test)
    print(f"checkGenIdxManual: {ak.all(delta_rs <= 0.4)}")
 
checkGenIdxRecommended(events)
checkGenIdxManual(events)

This gives me:

checkGenIdxRecommended: True
checkGenIdxManual: True

However strangely, if I separate the Jet and GenJet, I get a different result:

def checkGenIdxSeparate(jets, genjets):
    matched_genjets = genjets._apply_global_index(jets.genJetIdxG)
    delta_rs = jets.delta_r(matched_genjets)
    print(f"checkGenIdxOld: {ak.all(delta_rs <= 0.4)}")
jets = events.Jet.compute()
genjets = events.GenJet.compute()
checkGenIdxSeparate(jets, genjets)

The code above prints: checkGenIdxOld: False
I suppose the intended workflow is to keep the Jet and GenJet within the confines of events as long as possible.