Reading data with coffea NanoEvents

This is a rendered copy of nanoevents.ipynb. You can optionally run it interactively on binder at this link

NanoEvents is a Coffea utility to wrap flat nTuple structures (such as the CMS NanoAOD format) into a single awkward array with appropriate object methods (such as Lorentz vector methods\(^*\)), cross references, and nested objects, all lazily accessed\(^\dagger\) from the source ROOT TTree via uproot. The interpretation of the TTree data is configurable via schema objects, which are community-supplied for various source file types. These schema objects allow a richer interpretation of the file contents than the uproot.lazy methods. Currently available schemas include:

• BaseSchema, which provides a simple representation of the input TTree, where each branch is available verbatim as events.branch_name, effectively the same behavior as uproot.lazy. Any branches that uproot supports at “full speed” (i.e. that are fully split and either flat or single-jagged) can be read by this schema;

• NanoAODSchema, which is optimized to provide all methods and cross-references in CMS NanoAOD format;

• PFNanoAODSchema, which builds a double-jagged particle flow candidate colllection events.jet.constituents from compatible PFNanoAOD input files;

• TreeMakerSchema which is designed to read TTrees made by TreeMaker, an alternative CMS nTuplization format;

• PHYSLITESchema, for the ATLAS DAOD_PHYSLITE derivation, a compact centrally-produced data format similar to CMS NanoAOD; and

• DelphesSchema, for reading Delphes fast simulation nTuples.

We welcome contributions for new schemas, and can assist with the design of them.

\(^*\) Vector methods are currently made possible via the coffea vector methods mixin class structure. In a future version of coffea, they will instead be provided by the dedicated scikit-hep vector library, which provides a more rich feature set. The coffea vector methods predate the release of the vector library.

\(^\dagger\) Lazy access refers to only fetching the needed data from the (possibly remote) file when a sub-array is first accessed. The sub-array is then materialized and subsequent access of the sub-array uses a cached value in memory. As such, fully materializing a NanoEvents object may require a significant amount of memory.

In this demo, we will use NanoEvents to read a small CMS NanoAOD sample. The events object can be instantiated as follows:

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

NanoAODSchema.warn_missing_crossrefs = False

fname = "https://raw.githubusercontent.com/CoffeaTeam/coffea/master/tests/samples/nano_dy.root"
events = NanoEventsFactory.from_root(
    {fname: "Events"},
    schemaclass=NanoAODSchema,
    metadata={"dataset": "DYJets"},
).events()

In the factory constructor, we also pass the desired schema version (the latest version of NanoAOD can be built with schemaclass=NanoAODSchema) for this file and some extra metadata that we can later access with events.metadata. In a later example, we will show how to set up this metadata in coffea processors where the events object is pre-created for you. Consider looking at the from_root class method to see all optional arguments.

The events object is an awkward array, which at its top level is a record array with one record for each “collection”, where a collection is a grouping of fields (TBranches) based on the naming conventions of NanoAODSchema. For example, in the file we opened, the branches:

Generator_binvar
Generator_scalePDF
Generator_weight
Generator_x1
Generator_x2
Generator_xpdf1
Generator_xpdf2
Generator_id1
Generator_id2

are grouped into one sub-record named Generator which can be accessed using either getitem or getattr syntax, i.e. events["Generator"] or events.Generator. e.g.

events.Generator.id1.compute()

[1,
 -1,
 -1,
 21,
 21,
 4,
 2,
 -2,
 2,
 1,
 ...,
 1,
 -2,
 2,
 1,
 2,
 -2,
 -1,
 2,
 1]
--------------------------------------------------------------
type: 40 * int32[parameters={"__doc__": "id of first parton"}]

# all names can be listed with:
events.Generator.fields

['binvar', 'scalePDF', 'weight', 'x1', 'x2', 'xpdf1', 'xpdf2', 'id1', 'id2']

In CMS NanoAOD, each TBranch has a self-documenting help string embedded in the title field, which is carried into the NanoEvents, e.g. executing the following cell should produce a help pop-up:

Type:            Array
String form:     [1, -1, -1, 21, 21, 4, 2, -2, 2, 1, 3, 1, ... -1, -1, 1, -2, 2, 1, 2, -2, -1, 2, 1]
Length:          40
File:            ~/src/awkward-1.0/awkward1/highlevel.py
Docstring:       id of first parton
Class docstring: ...

where the Docstring shows information about the content of this array.

events.Generator.id1?

Type:            Array
String form:     dask.awkward<id1, npartitions=1>
File:            /opt/homebrew/lib/python3.11/site-packages/dask_awkward/lib/core.py
Docstring:       id of first parton
Class docstring:
Partitioned, lazy, and parallel Awkward Array Dask collection.

The class constructor is not intended for users. Instead use
factory functions like :py:func:`~dask_awkward.from_parquet`,
:py:func:`~dask_awkward.from_json`, etc.

Within dask-awkward the ``new_array_object`` factory function is
used for creating new instances.

Based on a collection’s name or contents, some collections acquire additional methods, which are extra features exposed by the code in the mixin classes of the coffea.nanoevents.methods modules. For example, although events.GenJet has the fields:

events.GenJet.fields

['eta', 'mass', 'phi', 'pt', 'partonFlavour', 'hadronFlavour']

we can access additional attributes associated to each generated jet by virtue of the fact that they can be interpreted as Lorentz vectors:

events.GenJet.energy.compute()

[[217, 670, 258],
 [34.5, 98.3, 1.16e+03, 38.1, 20.4, 29.7],
 [306, 62.8, 74.1, 769, 11.2],
 [170, 117, 29.3, 45.9],
 [101, 117, 129, 15.6],
 [63.1, 37.2, 33.7, 36.2],
 [303, 50.5, 1.29e+03, 278],
 [615, 282, 2.11e+03],
 [195, 47.6],
 [95, 44.6, 223, 318, 30, 108, 62.9],
 ...,
 [41.6, 36.7, 78.9, 13],
 [1.51e+03, 1.23e+03],
 [152, 160, 777, 27.1, 346, 65.1, 37.9, 27.2, 16.3],
 [35.4, 20.4],
 [20.1, 16.2],
 [34],
 [553, 283],
 [771, 452, 16],
 [76.9]]
----------------------------------------------------
type: 40 * var * float32

We can call more complex methods, like computing the distance \(\Delta R = \sqrt{\Delta \eta^2 + \Delta \phi ^2}\) between two LorentzVector objects:

# find distance between leading jet and all electrons in each event
dr = events.Jet[:, 0].delta_r(events.Electron)
dr.compute()

[[],
 [3.13],
 [3.45, 2.18],
 [1.58, 3.76],
 [],
 [0.053],
 [0.0748],
 [],
 [],
 [1.82],
 ...,
 [0.00115],
 [],
 [0.0149],
 [],
 [0.0308],
 [],
 [0.0858],
 [],
 []]
------------------------
type: 40 * var * float32

# find minimum distance
ak.min(dr, axis=1).compute()

[None,
 3.13,
 2.18,
 1.58,
 None,
 0.053,
 0.0748,
 None,
 None,
 1.82,
 ...,
 0.00115,
 None,
 0.0149,
 None,
 0.0308,
 None,
 0.0858,
 None,
 None]
-------------------
type: 40 * ?float32

# a convenience method for this operation on all jets is available
print(events.Jet.nearest(events.Electron).compute())

[[None, None, None, None, None], [Electron, Electron, ..., Electron], ..., [None, None]]

The assignment of methods classes to collections is done inside the schema object during the initial creation of the array, governed by the awkward array’s __record__ parameter and the associated behavior. See ak.behavior for a more detailed explanation of array behaviors.

Additional methods provide convenience functions for interpreting some branches, e.g. CMS NanoAOD packs several jet identification flag bits into a single integer, jetId. By implementing the bit-twiddling in the Jet mixin, the analsyis code becomes more clear:

print(events.Jet.jetId.compute())
print(events.Jet.isTight.compute())

[[6, 6, 6, 6, 6], [6, 2, 6, 6, 6, 6, 6, 0], ..., [6, 6, 0, ..., 6, 6], [6, 6]]
[[True, True, True, True, True], [True, True, ..., False], ..., [True, True]]

We can also define convenience functions to unpack and apply some mask to a set of flags, e.g. for generated particles:

print(f"Raw status flags: {events.GenPart.statusFlags.compute()}")
events.GenPart.hasFlags(['isPrompt', 'isLastCopy']);

Raw status flags: [[10625, 27009, 4481, 22913, 257, ..., 13884, 13884, 13884, 12876, 12876], ...]

CMS NanoAOD also contains pre-computed cross-references for some types of collections. For example, there is a TBranch Electron_genPartIdx which indexes the GenPart collection per event to give the matched generated particle, and -1 if no match is found. NanoEvents transforms these indices into an awkward indexed array pointing to the collection, so that one can directly access the matched particle using getattr syntax:

events.Electron.matched_gen.pdgId.compute()

[[],
 [-11],
 [-11, 11],
 [22, None],
 [],
 [None],
 [None],
 [],
 [],
 [11],
 ...,
 [11],
 [],
 [11],
 [],
 [-11],
 [],
 [None],
 [],
 []]
---------------------------------------------------------
type: 40 * var * ?int32[parameters={"__doc__": "PDG id"}]

events.Muon[ak.num(events.Muon)>0].matched_jet.pt.compute()

[[84.4, 29.4],
 [31.1],
 [53.4, 81.9],
 [29.2],
 [17.5],
 [65.9, 47.8],
 [58.5, 44.7],
 [50.2, 45.2],
 [33.3, 25.9],
 [None],
 [26.1],
 [25.8]]
-------------------------------------------------------
type: 12 * var * ?float32[parameters={"__doc__": "pt"}]

For generated particles, the parent index is similarly mapped:

events.GenPart.parent.pdgId.compute()

[[None, None, 1, 1, 23, 23, 23, 23, ..., 15, -15, -15, -15, -15, -15, 111, 111],
 [None, None, -1, 23, 23, 23, 23, ..., -11, None, None, None, None, None, 433],
 [None, None, -1, -1, 23, 23, 23, 23, ..., -423, -1, -1, -421, -421, 111, 111],
 [None, None, 21, 21, 23, -1, 23, 23, ..., -15, -15, -15, -15, -15, 111, 111],
 [None, None, 21, 21, 23, 23, 23, 23, ..., 13, 13, -13, 1, None, None, 2, 2],
 [None, None, 4, 23, 23, 23, 23, 23, ..., -15, -15, -15, 15, 15, 15, 423, 311],
 [None, None, 2, 2, 2, 23, 23, 2, 23, ..., -13, 13, 2, 2, 2, 2, 111, 111, 111],
 [None, None, -2, -2, 23, 21, 21, 23, 23, ..., 21, 21, 21, 21, None, 423, 2, 2],
 [None, None, 2, 23, 23, 23, 23, 23, ..., -15, -15, -15, -15, 111, 111, 311],
 [None, None, 1, 1, 1, 23, 21, 23, 23, 23, ..., -411, 21, 21, 1, 1, 1, 1, 3, 3],
 ...,
 [None, None, 1, 23, 23, 23, 23, 23, ..., -15, -15, -15, -15, 1, 1, 111, 111],
 [None, None, -2, 23, 23, 23, 23, 23, 13, 13, -13, -13, -13, 13],
 [None, None, 2, 2, 2, 23, 2, 23, ..., None, -413, 413, 413, 2, 2, -421, -421],
 [None, None, 1, 23, 23, 23, 23, 23, ..., -15, 15, 15, 15, -15, -15, -15, -15],
 [None, None, 2, 2, 23, 23, 21, 23, ..., 15, 15, 15, -15, -15, -15, 111, 111],
 [None, None, -2, 23, 23, None, 23, 23, ..., -15, -15, -15, 423, 4, 4, 3, 3],
 [None, None, -1, 23, 23, 23],
 [None, None, 2, 23, 23, 23, 23, -11, -11, 11],
 [None, None, 1, 1, 23, 23, 23, 23, ..., -15, -15, -15, 111, 111, 111, 111]]
--------------------------------------------------------------------------------
type: 40 * var * ?int32[parameters={"__doc__": "PDG id"}]

In addition, using the parent index, a helper method computes the inverse mapping, namely, children. As such, one can find particle siblings with:

events.GenPart.parent.children.pdgId.compute()
# notice this is a doubly-jagged array

[[None, None, [23, 21], ..., [-16, 111, ..., 211, -211], [22, 22], [22, 22]],
 [None, None, [23], [23], [23], [23], ..., None, None, None, None, None, [431]],
 [None, None, [23, -1], [23, -1], [23], ..., [13, -14], [13, -14], [22], [22]],
 [None, None, [23, -1], ..., [-16, 111, ..., 211, -211], [22, 22], [22, 22]],
 [None, None, [23, 1], [23, 1], [23], ..., None, None, [11, -11], [11, -11]],
 [None, None, [23], [23], ..., [16, 13, -14], [16, 13, -14], [421], [310]],
 [None, None, [23, 2, 2], [23, 2, 2], ..., [...], [22], [11, -11], [11, -11]],
 [None, None, [23, 21], [23, 21], [23], ..., None, [421], [11, -11], [11, -11]],
 [None, None, [23], [23], ..., [-16, 111, ..., 311], [22, 22], [22, 22], [310]],
 [None, None, [23, 21, 21], [23, ...], ..., [11, -11], [11, -11], [11, -11]],
 ...,
 [None, None, [23], [23], [23], ..., [11, -11], [11, -11], [22, 22], [22, 22]],
 [None, None, [23], [23], [23], ..., [...], [-13], [-13, 22], [-13, 22], [13]],
 [None, None, [23, 2, 21], ..., [2, 21, ..., 11, -11], [13, -14], [13, -14]],
 [None, None, [23], ..., [...], [-16, 211, 211, -211], [-16, 211, 211, -211]],
 [None, None, [23, 21], [23, 21], ..., [-16, -11, 12], [22, 22], [22, 22]],
 [None, None, [23], [23], ..., [423, -421, 11, -11], [11, -11], [11, -11]],
 [None, None, [23], [23], [-13, 13], [-13, 13]],
 [None, None, [23], [23], [23], ..., [-11, 11], [-11, 22], [-11, 22], [11]],
 [None, None, [23, 21], [23, 21], ..., [22, ...], [22, 22], [22, 22], [22, 22]]]
--------------------------------------------------------------------------------
type: 40 * var * option[var * ?int32[parameters={"__doc__": "PDG id"}]]

Since often one wants to shortcut repeated particles in a decay sequence, a helper method distinctParent is also available. Here we use it to find the parent particle ID for all prompt electrons:

events.GenPart[
    (abs(events.GenPart.pdgId) == 11)
    & events.GenPart.hasFlags(['isPrompt', 'isLastCopy'])
].distinctParent.pdgId.compute()

[[],
 [23, 23],
 [23, 23],
 [],
 [],
 [],
 [],
 [23, 23],
 [],
 [23, 23],
 ...,
 [],
 [],
 [23, 23],
 [],
 [],
 [],
 [],
 [23, 23],
 []]
---------------------------------------------------------
type: 40 * var * ?int32[parameters={"__doc__": "PDG id"}]

Events can be filtered like any other awkward array using boolean fancy-indexing

mmevents = events[ak.num(events.Muon) == 2]
zmm = mmevents.Muon[:, 0] + mmevents.Muon[:, 1]
zmm.mass.compute()

[94.6,
 87.6,
 88,
 90.4,
 89.1,
 31.6]
-----------------
type: 6 * float32

# a convenience method is available to sum vectors along an axis:
mmevents.Muon.sum(axis=1).mass.compute()

[94.6,
 87.6,
 88,
 90.4,
 89.1,
 31.6]
-----------------
type: 6 * float32

As expected for this sample, most of the dimuon events have a pair invariant mass close to that of a Z boson. But what about the last event? Let’s take a look at the generator information:

print(mmevents[-1].Muon.matched_gen.pdgId.compute())
print(mmevents[-1].Muon.matched_gen.hasFlags(["isPrompt"]).compute())

[-13, 13]
[False, False]

So they are real generated muons, but they are not prompt (i.e. from the initial decay of a heavy resonance)

Let’s look at their parent particles:

mmevents[-1].Muon.matched_gen.parent.pdgId.compute()

[-15,
 15]
--------------------------------------------------
type: 2 * ?int32[parameters={"__doc__": "PDG id"}]

aha! They are muons coming from tau lepton decays, and hence a fair amount of the Z mass is carried away by the neutrinos:

print(mmevents[-1].Muon.matched_gen.sum().mass.compute())
print(mmevents[-1].Muon.matched_gen.parent.sum().mass.compute())

31.265245
91.68365

One can assign new variables to the arrays, with some caveats:

• Assignment must use setitem (events["path", "to", "name"] = value)

• Assignment to a sliced events won’t be accessible from the original variable

• New variables are not visible from cross-references

mmevents["Electron", "myvariable"] = mmevents.Electron.pt + zmm.mass
mmevents.Electron.myvariable.compute()

[[],
 [121],
 [],
 [],
 [],
 []]
-----------------------
type: 6 * var * float32