Learning the awkward arrays syntax

The core of coffea leverages a completely new syntax for expressing analysis computations: awkward arrays and its index-based notation. For people coming from a more traditional loop-based programming syntax, the syntax will take some getting use to, but this tutorial can hopefully help you understand how to understand the syntax and how to understand the various method.

Let use begin by first understanding what you need to explore the contents of a typical n-tuple file using coffea related tools. For the examples below we will be using the CMS open data. Either download the file to your working directory to follow along.

First import the relevant coffea objects:

from coffea.nanoevents import NanoEventsFactory
from coffea.nanoevents.schemas import  NanoAODSchema
import numpy as np
import awkward1 as ak

Now we can create the event list as an awkward array using coffea tools like:

events = NanoEventsFactory.from_root(
    'file:TTbar.root',    # The file, notice the prefix "file:" for local file operation
    'Events',             # Name of the tree object to open
    entry_stop=50,        # Limit the number of events to process, nice for small scale debugging
    schemaclass=NanoAODSchema
).events()

The last schemaclass argument will be let unexplained for now, see the schema tutorials to learn more about what this is. Here we have created the events as an awkward array. To see that is stored in the array we can use:

print(events.fields)

['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event']

Indicating the collections that are stored in the awkward array. To see how many events exists in our file we can use the typical python method:

print(len(events))

50

The 50 here in corresponds correctly to the entry_stop used in to open the file. Next we can, of course, start to explore the contents of the various object collections. One can access the fields of the event as if it was a regular data member

print(events.Muon.fields)
print(events.Jet.fields)

['pt', 'eta', 'phi', 'mass', 'charge', 'pfRelIso03_all', 'pfRelIso04_all', 'tightId', 'softId', 'dxy', 'dxyErr', 'dz', 'dzErr', 'jetIdx', 'genPartIdx', 'jetIdxG', 'genPartIdxG']
['pt', 'eta', 'phi', 'mass', 'puId', 'btag']

Ah-ha! Now, we are starting to see numbers we can play around with. Notice that coffea was written with high energy physics analysis in mind, so even if the electron energy doesn’t look like it is stored from the output of the fields, we can still access methods that we will typically associate with 4-vectors. In particular, notice that we can call the energy field of the electron collection, even though the energy field isn’t explicitly defined. coffea is designed with 4 vectors in mind, so the energy collection is calculated on the fly.

print(events.Muon.pt)
print(events.Muon.energy)

[[23.6], [], [], [57.3, 29.3, 7.7], [], [], ... [], [], [57.2], [], [19.7], [], []]
[[52.3], [], [], [95.6, 30.2, 8.01], [], [], ... [], [], [309], [], [19.7], [], []]

Now, looking at the output, we can begin to get a grasp of what awkward arrays are: the variable events.Muon.pt variable represents an \(N\) events times \(A\) objects array of floating point, the events.Muon variable represents and \(N\) events times \(A\) objects, acting as a collection of floating point arrays.

The “awkward” part of awkward array refers to two parts: - The value of A is different for each event and for each collection. In this demonstration, our 0-th event has 1 muon, the 1-st and 2-rd event has no muons, and the 3-rd event has 3 muons and so on. - Each collection can have a different number of fields. In a sense, the events, Muon and pt variables are just an easy way for representing the final NxA array that we might be interested in for the analysis.

In our case the N number of events is what will refer to as the outermost dimension or axis of the various objects, A is the one inner dimension of the of array. One thing to keep in mind is that we typically refer to the dimension of events.Muon and events.Muon.pt as both being NxA, except events.Muon is a NxA array of objects (or collection), while the events.Electron.pt being a NxA array of a concrete data type (in this case a float).

We can use the usual index notation to look at a particular object of interest. For instance if we want to look at the 0-th Muon of the 3-rd event in our event list, we can write:

print(events[3].Muon[0].pt, events[3].Muon[0].eta)

57.327919006347656 1.0991891622543335

But the real power of using awkward arrays is for using awkward arrays comes in when you don’t explicitly use a concrete index, and instead call calculation all in an abstract form

Basic object and event selection

Let us start with the most basic example of event selection. Say we want to select event with electrons that have \(p_T > 30\) GeV and \(|\eta| < 2.0\). The awkward array allows us to write something like:

mask_pt = events.Muon.pt > 50
mask_eta = np.abs(events.Muon.eta) < 2.0
muon_mask = mask_pt & mask_eta
print(mask_pt)
print(mask_eta)
print(muon_mask)

[[False], [], [], [True, False, False], [], ... [], [True], [], [False], [], []]
[[True], [], [], [True, True, True], [], [], ... [], [], [False], [], [True], [], []]
[[False], [], [], [True, False, False], [], ... [], [False], [], [False], [], []]

We can see that the usual logic comparison operators generate a NxA boolean array telling use which electron (or more specifically which electron.pt and electron etas) pass this particular selection criteria. This particular boolean array generated from a logic operation on usual arrays is typically call a mask. We can use the typical boolean operation & operation to get the intersect of multiple masks, or maybe the | operator for the union. Now the problem is where can we use this mask? The answer is any array that has a NxA structure and can receive these masks to create a reduced array!

print(events.Muon.pt[muon_mask])
print(events.Muon.eta[muon_mask])
selectedMuons = events.Muon[muon_mask]
print(selectedMuons.pt)

[[], [], [], [57.3], [], [], [], [], [], ... [], [], [], [], [], [], [], [], []]
[[], [], [], [1.1], [], [], [], [], [], ... [], [], [], [], [], [], [], [], []]
[[], [], [], [57.3], [], [], [], [], [], ... [], [], [], [], [], [], [], [], []]

Probably the most important place to put the mask is directly in the events.Muon index, this generates a new collection of electrons that preserves the NxA structure, but have various collection instances filtered out. If you are familiar with numpy, this sort of index-based array filtering look familiar. The difference is that because awkward arrays accept arrays of varying inner dimensions, it can truly preserve the structure of such selection, rather than having everything be flattened out.

x = np.array([1,2,3,4,5,6,7,8,1,1,1,2])
print( x[x% 2 == 0])
y = np.array([[1,2,3,4],[5,6,7,8],[1,1,1,2]])
print( y[y%2==0])
z = ak.Array([[1,2,3,4],[5,6,7,8],[1,1,1,2]])
print(z[z%2==0])

[2 4 6 8 2]
[2 4 6 8 2]
[[2, 4], [6, 8], [2]]

Now suppose we only want events that have at least 1 electron selected event. What we need are a set of functions that can reduces this NxA' array to something of just dimension N. Formally this is called reduction operations, and the awkward package has a large set of functions that can reduce the dimension of arrays. In our case, what we want is:

muon_count = ak.count(selectedMuons.pt, axis=-1)
event_mask = muon_count >= 1
print(event_mask)

[False, False, False, True, False, False, ... False, False, False, False, False]

To break this down, ak.count, as the method name suggests “counts” the number of elements along a certain axis, in our case, what we are interested is the innermost dimension/axis, hence the typical python notation of axis=-1. Using this we can run the event selection using the usual masking notation:

selectedEvents = events[event_mask]
print(event_mask)
print(events.Muon.pt)
print(selectedEvents.Muon.pt)
print(len(selectedEvents))

[False, False, False, True, False, False, ... False, False, False, False, False]
[[23.6], [], [], [57.3, 29.3, 7.7], [], [], ... [], [], [57.2], [], [19.7], [], []]
[[57.3, 29.3, 7.7], [52.2]]
2

Here we can confirm that the first event to pass the event selection is the 1-st event in the event list, and the 0-th instance in the selectedEvents.Muon.pt result of the selectedEvents indeed corresponds to the values stored in the 1-st event of the original event list.

Object storage and collection creation

Having completed the selection, we might be rather annoyed that we didn’t just store the selected Electron, since these are the objects that we are likely going to use for further calculation. Following from the code above, what we can do is add the additional selection to the selectedMuon collections. This is valid since the N dimensional event mask “makes sense” performed on the NxA' dimensional selectedMuon object.

our_selectedMuons = selectedMuons[event_mask]
print(our_selectedMuons.pt)
print(len(our_selectedMuons))

[[57.3], [52.2]]
2

However, this is rather undesirable, since now we have some a bunch of detected collections, and event lists that we need to take care of: selectedMuon, selectedEvents, our_selectedMuons. And this is with just one toy object selection. One can imagine if there isn’t some sort of way to store collections into events, the analysis code will get out of hands very quick. This also ties into the topic that there might be certain physics quantities that are specific to a certain analysis that might be used for the analysis object selection and would be nice to add to the electron collection if it isn’t a standard variable that is maintained by the NanoAOD development team. Here we are going to add a very artificial example of calculating the inverse of the electron pt, then selecting on the inverse pt. This very simple example will demonstrate the typical syntax used for storing variables as well as exposing one of the particuliar quirks of awkward arrays:

print('First attempt at adding extended variables to events')
events.Muon['invpt'] = 1/events.Muon.pt
events['selectedMuon_1'] = events.Muon[events.Muon.pt > 50]

print(events.fields)
print(events.Muon.fields)
print(events.selectedMuon_1.fields)

print('\n\nSecond attemp at adding extended variables to events')
events['myMuon'] = events.Muon[:]
events.myMuon['invpt'] = 1/events.myMuon.pt
events['selectedMuon_2'] = events.myMuon[events.myMuon.pt > 50]

print(events.fields)
print(events.myMuon.fields)
print(events.selectedMuon_2.fields)

print('\n\nThird attemp at adding extended variables to events')
myMuon = events.Muon[:]
myMuon['invpt'] = 1/myMuon.pt
events['selectedMuon_3'] = myMuon[myMuon.pt > 50]

print(events.fields)
print(myMuon.fields)
print(events.selectedMuon_3.fields)

First attempt at adding extended variables to events
['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']


Second attemp at adding extended variables to events
['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']


Third attemp at adding extended variables to events
['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']

Let’s get the straightforward part of the code clear up. The addition of collections looks very straight forward, one can think of the events as something that looks like a “dictionary of collection with a common outer dimension”, so the addition of the two electron collections to the event has a very distionary-esque notation. What is strange is the persistence of the extended collection for the electrons. Logically, the operation looks identical, but the first attempt to add the new variable invpt directly to events.Muon fails to persist, and thus all direct extensions of events.Muon doesn’t include the new invpt field.

The reason for this is rather technical regarding the mutability of objects in python and awkward. The rule-of-thumb is that collections that are directly generated from the file, (a.k.a. the collections directly obtained listed the events.fields immediate after opening a file) can never be altered, and therefore cannot have extended variables added to them. To create an extended variable to some collection, we will need to make some sort of copy of the original either by some trivial kinematic selection (ex. myMuons = events.Muon[events.Muon.pt > 0]) or some trivial splicing (myMuon = events.Muons[:]). Another feature of mutability is that once the collection is added to the event collection, it becomes immutable. That is why the third attempt is the one that adds the both the electron extended variable and the extended electron collection to the event.

Because of these quirks, it would typically be worth it to wrap the object selection into a function if the object selection is typical within an analysis, and it also helps with code readability

def SelectMuons(muon):
    muon = muon[muon.pt > 50]
    muon['invpt']  = 1.0 / muon.pt
    return muon

events['selectedMuon_f'] = SelectMuons(events.Muon)
print(events.fields)
print(events.selectedMuon_f.fields)

['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']

Once the new object collection has been added to the event collection, they will persist to arbitrary levels of event selection:

myevents = events[ak.count(events.selectedMuon_f.pt,axis=-1) > 0 ]

print(myevents.fields)
print(myevents.selectedMuon_f.fields)

myevents = events[ak.count(events.selectedMuon_f.pt,axis=-1) > 1 ]

print(myevents.fields)
print(myevents.selectedMuon_f.fields)
myevents = events[ak.count(events.selectedMuon_f.pt,axis=-1) > 2 ]

print(myevents.fields)
print(myevents.selectedMuon_f.fields)

['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']
['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']
['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']
['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']

Comparing array based computations and loop based computation

So to put this together into a single code block, suppose our analysis consisted of selecting events that have at least 1 electron with \(p_{T} > 50GeV\), \(|\eta| < 2.4\), and we want to calculate the average of all such electron’s inverse \(p_{T}\) within the selected events. Our awkward array code would look something like:

events = NanoEventsFactory.from_root( 'file:TTbar.root',
                                      'Events',
                                      entry_stop=50,
                                      schemaclass=NanoAODSchema).events()

## Object selection
selectedMuon = events.Muon[ (events.Muon.pt > 50) &
                            (np.abs(events.Muon.eta)<2.4) ]
selectedMuon['invpt'] = 1/selectedMuon.pt
events['selectedMuon'] = selectedMuon

# Event selection
events = events[ak.count(events.selectedMuon.pt,axis=-1) >= 1]

# Calculating the total average
print(ak.sum(events.selectedMuon.invpt)/ak.count(events.selectedMuon.invpt))

0.018032373239596684

In total, this is 4 statements (not counting the file reading step) used to make this analysis. Compare that with the loop based notation:

events = NanoEventsFactory.from_root( 'file:TTbar.root',
                                      'Events',
                                      entry_stop=50,
                                      schemaclass=NanoAODSchema).events()

count = 0
suminv = 0
for i in range(len(events)):
    is_good = []
    for j in range(len(events[i].Muon)):
        if events[i].Muon[j].pt > 50 and np.abs(events[i].Muon[j].eta) < 2.4:
            is_good.append(j)
    if len(is_good) >= 1:
        for j in is_good:
            count = count +1
            suminv += 1.0/ events[i].Muon[j].pt

print(suminv/count)

0.01803237348586199

Notice the results are only difference because the 32-bit to 64-bit float conversion is happening at different places. For awkward arrays, this is happening only after the sum has been performed. For the loop based approach this happening every time the += operator is called.

In the loop based analysis, notice for such a simple analysis, many lines of code are dedicated to just bookkeeping stuff: number of electrons passing criteria, adding a counter variable and sum variable… etc., instead of actually analysis computation. The array based notation for expressing the analysis is much cleaner, if rather more unfamiliar to typical users.

Of course, this isn’t the end. Physics analysis are typically more involved than just basic selection and counting. In the next session, we will talk about how to perform more involved calculations with awkward arrays that use multiple collections within an event collection.

More on index notation

Let us dive in a bit more regarding the index notation, with a fresh set of events.

events = NanoEventsFactory.from_root( 'file:TTbar.root',
                                      'Events',
                                      entry_stop=1000,
                                      schemaclass=NanoAODSchema).events()

The index notation largely follows in index notation of numpy arrays, which means we can use some of the following

• array[ -n ]: to the objects in the array in reverse order

• array[ start:end:step ]: to slice an array given a certain range and step size,

• array[ [list, of, indexes] ]: create a new array composed of the given indices

This is very simply when the array is just N dimensional (a.k.a. contains a single index). Let us use the most simple example, the events.MET.pt object is just a list floats, with the length matching the number of events. Using this we can get:

print("The entire MET list:                  ", events.MET.pt)
print("The MET from the second of last event:", events.MET.pt[-2])
print("The MET of the first 5 events:        ", events.MET.pt[0:5])
print("The MET of every other event:         ", events.MET.pt[0::2])
print("The MET of events 1 3 and 10:         ", events.MET.pt[[1,3,10]])

The entire MET list:                   [18.3, 50.1, 6.06, 38.8, 4.99, 68, 51.1, ... 8.04, 30.6, 107, 20.8, 31.5, 85.6, 40.6]
The MET from the second of last event: 85.59586334228516
The MET of the first 5 events:         [18.3, 50.1, 6.06, 38.8, 4.99]
The MET of every other event:          [18.3, 6.06, 4.99, 51.1, 28.5, 12.7, 16.8, ... 56.2, 87.7, 133, 30.6, 20.8, 85.6]
The MET of events 1 3 and 10:          [50.1, 38.8, 12.7]

Very rarely will a typical analysis want to directly operator on the event index. What we can do is operate on object indices using the notation:

object[event_index, object_index]

There the comma in the square braces is used to separate the first index argument that operation on the outermost dimension (in our case the event index), while the second argument is for the inner dimension index (in which our case the object index within an event). For instance, if we want to get the leading jet’s \(p_{T}\) for every event, we can use the following

print(events.Jet.pt[:,0])

[147, 76.9, 84.4, 52.8, 105, 94.3, 34.8, ... 41.7, 133, 210, 91.3, 76.3, 96.2, 139]

The first argument : is a trivial slice operation, indicating that we want every event. The second operator indicates that we want the leading (a.k.a. 0-th) jet in each event. Notice that the user is responsible to make sure a corresponding object exists. For instance, if we run the same notation using the muon collection instead:

print(events.Muon.pt[:,0])

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-22-af27855176d3> in <module>
----> 1 print(events.Muon.pt[:,0])

~/VirtualENV/coffeaenv/lib/python3.8/site-packages/awkward1/highlevel.py in __getitem__(self, where)
    962         have the same dimension as the array being indexed.
    963         """
--> 964         return awkward1._util.wrap(self._layout[where], self._behavior)
    965
    966     def __setitem__(self, where, what):

ValueError: in ListArray64 attempting to get 0, index out of range

(https://github.com/scikit-hep/awkward-1.0/blob/0.4.5/src/cpu-kernels/awkward_NumpyArray_getitem_next_at.cpp#L21)

We get an out of range error, since not every event has a muon. In this case, we would first need to filter on events to ensure events contain muons, and then run same notation

events_with_muons = events[ak.count(events.Muon.pt,axis=-1) > 0 ]
print(events_with_muons.Muon.pt[:,0])

[23.6, 57.3, 25.6, 3.13, 12.8, 17.9, 8.44, ... 3.42, 7.64, 8.26, 3.29, 27, 14.6, 3.3]

The second argument doesn’t need to be singular either, for instance, if our analysis is only interested in the 2 leading jets within an event, what we can do is:

events_with_2jets = events[ak.count(events.Jet.pt,axis=-1) >= 2 ]
events_with_2jets['selectedJetsPt'] = events_with_2jets.Jet[:,0:2]
print(events_with_2jets.selectedJetsPt.pt)

[[147, 72.7], [76.9, 70.5], [84.4, 60.5], ... [76.3, 72.7], [96.2, 77.5], [139, 106]]

Where we have now created the collection which contains just the 2 leading jets in the event using a slicing notation. The list of index notation is also helpful if we want to reorder objects. For example if we want to find the 2 jets that are the most forward in the detector, we can use awkward1 methods to generate the required index

order = ak.argsort(np.abs(events.Jet.eta), axis=-1 , ascending=False)
print(events.Jet.eta)
print(order)
events['JetByEta'] = events.Jet[order]
print(events.JetByEta.eta)

[[-2.1, -1.46, -1.73, -1.57, 0.0878], ... 0.94, -1.85, 3.86, 4.29, -1.51, -2.49]]
[[0, 2, 3, 1, 4], [1, 4, 0, 2, 3], [5, ... 2, 0, 5, 1, 3], [5, 4, 7, 3, 1, 6, 2, 0]]
[[-2.1, -1.73, -1.57, -1.46, 0.0878], ... -2.49, -1.85, -1.75, -1.51, 0.94, 0.0518]]

Now this example is a bit more involved:

• First the method `ak.argsort <https://awkward-array.readthedocs.io/en/latest/_auto/ak.argsort.html>`__, generates the index by which the object should be called along the last axis(axis=-1) in descending order: For example in the first event, the jet with the largest \(|\eta|\) is the 0-th on, followed by the 2-nd one and so on. The order variable would be a NxA array of integers.

• Next, when placed in the square parenthesis, since the outer dimension of the order matches the outer dimension of events.Jet it automatically knows to apply the index by event to each inner dimension.

Notice because of how the second line is phrased, the resulting events.JetByEta has an inner dimension that is the same as the order variable and not the events.Jet variable. For instance, if we are only interested in the 2 least forward jets we can use:

new_order = order[:, -2:]
events['JetByEta'] = events.Jet[new_order]
print(events.JetByEta.eta)

[[-1.46, 0.0878], [-0.27, -0.178], [0.738, ... [-0.399, -0.239], [0.94, 0.0518]]

Another handy application of this is the index notation is the old process of using object indices in the place of pointers. In this NanoAOD event, muons have a genPartIdx field, which indicates which events.GenPart collection it should be matched to. Notice that we will need to some additional parsing to ensure, that the muons of interest have correct genPartIdx values (>= 0)

muon_withgen = events.Muon[events.Muon.genPartIdx >= 0]
muon_genpart = events.GenPart[muon_withgen.genPartIdx]
print(muon_withgen.pt)
print(muon_genpart.pt)

[[23.6], [], [], [57.3, 29.3, 7.7, ... [3.3, 4.25, 3.28, 19.6, 17.3, 6.62, 3.72]]
[[24], [], [], [56.2, 29.4, 29.4], ... [14.7], [3.31, 27, 27, 27, 27, 27, 3.31]]

Now since the muon_genpart has the same dimension as that of the events.MuonWithGen collection, we can add the muon_genpart to the MuonWithGen collection. This makes additional codes much more useful:

muon_withgen['GenPart'] = muon_genpart
print(muon_withgen.GenPart.pdgId)
print(muon_withgen.GenPart.status)

[[-13], [], [], [13, 13, 13], [], ... -13], [-13], [-13, 11, 11, 11, 11, 11, -13]]
[[1], [], [], [1, 1, 1], [], [], [2], ... [], [1, 1], [1], [1, 1, 1, 1, 1, 1, 1]]

Now the code mimics quite a bit with what you would normally get with object based code, which makes awkward arrays rather more useful than just the flat arrays one typically get in C++ n-tuple codes.

Calculation with objects

Up til now, the calculations have all removed around single collection. Even in the case of the cell involving two collection, one is just used to generate index dimension. The first instance of calculation, where we use a mathematic operation to generate a new array is in the case of mask generation and the calculation of the inverse of muon pt, but that apparently only involved 1 array. In this section, we dive a bit deeper into what mathematical operators actually imply in awkward.

The mathematical operations of awkward arrays is member-wise: Given 2 arrays of identical dimension (\(a\) and \(b\)) and a generic binary operation \(\otimes\), the statement \(a\otimes b\), will generate a new array with the same dimensions as \(a\), with output being a member-wise results of \(o_{ijk\ldots} = a_{ijk\ldots} \otimes b{ijk\ldots}\). Now what if \(a\) and \(b\) are different dimension, say in our case \(a\) is a scalar 1, and \(b\) is the muon \(p_{T}\) array. What happens is formally known as broadcasting: you can think that awkward automatically expands the scalar \(a\) into a NxA array and performs a typical member-wise arithmetics. Broadcasting is a rather complicated subject to fully discuss, but a rule-of-thumb is that the follow operations would be considered valid:

• scalar \(\otimes\) N: scalar is expanded to a N

• scalar \(\otimes\) NxA: scalar is expanded to a NxA array

• NxA \(\otimes\) NxA: typical member-wise calculation

• N \(\otimes\) N: expansion of the first N array to a NxA array

Since in awkward array, the A dimension is a variable one, we will never encounter the case of a A\(\otimes\)NxA array. To use some concrete examples:

print('scalar @ N   :', 100 +  events.MET.px)
print('scalar @ NxA :', 100 + events.Muon.px)
print('NxA    @ NxA :', events.Muon.px + events.Muon.py)
print('N      @ NxA :', events.MET.px + events.Muon.px)

scalar @ N   : [92.7, 145, 99.4, 121, 101, 35.7, 146, ... 108, 99.9, 207, 120, 76.7, 14.7, 139]
scalar @ NxA : [[89.8], [], [], [148, 116, 104], ... [88.3], [103, 104, 103, 119, 117, 107, 104]]
NxA    @ NxA : [[11.1], [], [], [16, -7.89, ... [3.99, 4.12, 3.91, 23.5, 21.4, 6.87, 4.26]]
N      @ NxA : [[-17.5], [], [], [68.7, 37.4, 25, ... [42.3, 43.3, 42.3, 58.2, 55.6, 45.7, 42.7]]

Pretty much all calculations can be broken down into binary operations like above, coffea even provides some common HEP function such as delta_r, so you don’t have to re-write common routines from scratch. The difficulty comes in mangling the array into the correct dimension, so that broadcasting can take over. For instance, the calculation of the \(\Delta R\) of the first and second leading jets in an event is simple, this is just a reduction of the 2 NxA array to 2 N array via index selection, then calling an in-built binary operation.

evt = events[ak.count(events.Jet.pt,axis=-1) >= 2 ]
first_jet = evt.Jet[:,0]
second_jet = evt.Jet[:,1]
print(first_jet.delta_r(second_jet))

[3.18, 4.35, 5.69, 0.947, 2.7, 1.8, 1.65, ... 2.9, 3.19, 2.89, 3.04, 2.85, 3.61]

To calculated of \(\Delta R\) the leading jet with all muon in the event is also simple (notice how it takes on the dimension of the Muon collection)

print(first_jet.delta_r(evt.Muon))

[[3.16], [], [], [2.66, 3.49, 3.56, ... [3.5, 3.47, 3.58, 3.55, 3.48, 3.55, 3.53]]

Now suppose we want to perform a jet-cone matching of the gen particles collection with the all jets in the event. How do we perform broadcasting of a NxA array with a NxB array? The answer is to expand one of the array into a NxAxB array for the broadcast for higher dimensions can be performed. Details of how this method works can be found the awkward documentation for the `cartesian <https://awkward-array.readthedocs.io/en/latest/_auto/ak.cartesian.html>`__ and `unzip <https://awkward-array.readthedocs.io/en/latest/_auto/ak.unzip.html>`__ method, but here is skeleton code performing object x object computations:

jets = evt.Jet[:]
_, genpart = ak.unzip(ak.cartesian([evt.Jet,evt.GenPart],nested=True))
jets['genpart'] = genpart[jets.delta_r(genpart) < 0.4]
print(jets.genpart.pt)

[[[], [24, 24, 24, 24], [], [], []], ... 27], [], [], [], [], [0.545, 0.545], []]]

Breaking down the code above:

The ak.unzip and ak.cartesian method expands the NxA evt.Jet array and the NxA evt.GenPart array in to a NxBxA arrays (named _ since we wouldn’t be using this) and NxAxB array genpart. The broadcast is the performed on the NxA jets array and NxAxB genpart array to generate a delta R mask of dimension NxAxB, which is then used to filter the genpart. Now we can store the NxAxB array into the NxA jet collection, which is used to find all gen particles associated with a jet via delta-R matching!

This is probably the most advanced that calculation that a typical analysis will need in computation: the computation of arrays of mismatch dimensions. Hopefully this introduction is enough to help people start to translate their loop based analysis calculation into array-based syntax.

When all else fails - numba for loops

TO BE WRITTEN: assigned -Yi-Mu Chen