CoffeaTeam / coffea

Basic tools and wrappers for enabling not-too-alien syntax when running columnar Collider HEP analysis.
https://coffea-hep.readthedocs.io
BSD 3-Clause "New" or "Revised" License
134 stars 128 forks source link

Issue with accessing children with NanoEvents after .compute() #1200

Open rkansal47 opened 3 weeks ago

rkansal47 commented 3 weeks ago

Describe the bug

If I try to access the children of an object loaded as a dask_awkward array and then computed I get the error:

File ~/mambaforge/envs/python311/lib/python3.11/site-packages/coffea/nanoevents/methods/nanoaod.py:139, in GenParticle.children(self)
    133 @dask_property
    134 def children(self):
    135     """
    136     Accessor to direct children of this particle (not grandchildren). Includes particles
    137     with the same PDG ID as this particle.
    138     """
--> 139     return self._events().GenPart._apply_global_index(self.childrenIdxG)

File ~/mambaforge/envs/python311/lib/python3.11/site-packages/coffea/nanoevents/methods/base.py:270, in NanoCollection._events(self)
    268 if "@original_array" in self.attrs:
    269     return self.attrs["@original_array"]
--> 270 return self.attrs["@events_factory"].events()

KeyError: '@events_factory'

with coffea 2024.10.0 awkward 2.6.9.

To Reproduce MRE:

events_delayed = nanoevents.NanoEventsFactory.from_root(
    {
        "root://cmseos.fnal.gov///store/user/lpcpfnano/cmantill/v2_3/2017/HH/GluGluToHHTobbVV_node_cHHH0_TuneCP5_13TeV-powheg-pythia8/GluGluToHHTobbVV_node_cHHH0/220808_163755/0000/nano_mc2017_1-1.root": "Events"
    },
    schemaclass=nanoevents.NanoAODSchema,
    delayed=True,
).events()

higgs = events_delayed.GenPart[events_delayed.GenPart.hasFlags(["fromHardProcess", "isLastCopy"]) * (events_delayed.GenPart.pdgId == 25)]
print(higgs.compute().children)
lgray commented 3 weeks ago

The issue here is that the self-reference to the original events is lost when you compute (because you can't recursively compute a dask array!). So this isn't actually a bug but rather a consequence of how dask works.

I would suggest instead using delayed=False and doing your data exploration with completely eager nanoevents on that single file or you can try higgs.children.compute() and then work with that array (which you will not be able to recurse through.

Or if you know what you're looking for you can work with the dask array itself and then only compute what you want.

rkansal47 commented 3 weeks ago

I see. For context, I was trying to compute earlier than needed as a workaround for #1199, i.e., trying to flatten higgs.children.compute().children instead of higgs.children.children which gives an error.