Introduction¶

This notebook is a quick starter guide on Uproot, Awkward Array, and Vector, which are all Python packages that are part of the Scikit-HEP project. We'll briefly go through parts of each and see how they work together.

Most of this notebook is adapted from my Uproot tutorial aimed at beginners in data analysis in Python. Here, I'm assuming that the reader is already familiar with at least Python, NumPy, and Matplotlib, so I've compressed the tutorial and added a bit of more advanced material, but it still just scratches the surface of these packages. See the links at the end for more extensive resources.

Motivation¶

Python¶

For quite a while now, ROOT has been the main software used by particle physicists for data analysis. ROOT is written in C++, and at least until recently most analysis code has also been written in C++. Outside of HEP, Python is a more commonly used language for data analysis, which means that there are many popular and well-supported tools that can be utilized for common tasks. Python is also generally considered easier to learn than C++.

Uproot¶

Uproot is a Python module that can read and write files in the ROOT format. It’s part of the Scikit-HEP project, a collection of Python tools for HEP data analysis. Uproot isn’t the only way to read/write ROOT files in Python; there are other packages that allow you to do this. However, there are some significant advantages of Uproot. We’ll go through a few here.

The only software required by Uproot (other than Python itself, of course) is NumPy. Awkward Array and Vector are not required but highly recommended for use with Uproot. Uproot, Awkward Array, and Vector are all part of the Scikit-HEP project. All of the above are Python modules, which can be installed with nothing more than pip or conda.

Most importantly, Uproot does not require any part of the ROOT software, which is quite large and can be non-trivial to install properly.

Uproot was designed for efficient reading of files. It uses compiled functions with vector optimizations to perform very fast calculations on arrays.

When you import a Python package that uses ROOT, there can be a delay of several seconds while all of the required libraries are loaded. This is required every time you restart the Python interpreter or run a script. By virtue of not using ROOT itself, Uproot doesn’t have this issue.

Awkward¶

One issue with manipulating HEP data in Python is that it is often not rectangular. That is, the length of object collections is not constant. For example, different events will have different numbers of jets, and different jets will have different numbers of associated tracks. Such arrays are called jagged arrays. Packages like NumPy are only designed to work with rectangular arrays, where each dimension has a fixed length throughout the entire array.

Jagged vs. rectangular arrays

The name “jagged array” comes from the idea that if you have a 2D array in which there’s a different number of entries in each row and you line rows up by their first element, then the right edge of the array will be “jagged” as opposed to straight:

Rectangular array:

[[0, 1],
 [2, 3],
 [4, 5],
 [6, 7],
 [8, 9]]

Jagged array:

[[0, 1],
 [2, 3, 4],
 [5],
 [6, 7],
 [8, 9]]

Awkward Array is a package used to deal with jagged arrays so that they can be manipulated in ways almost identical to standard rectangular NumPy arrays. Other Python packages for reading ROOT files usually do not support jagged arrays, which means that you cannot do operations on multiple events at once if you have jagged data. But Uproot and Awkward Array can, and this is generally much faster than operating on one event at a time.

Vector¶

Vector is a package that provides an interface for 2D, 3D, and (3+1)D vectors and can be used with NumPy and Awkward Array to perform operations on arrays of vectors. This includes calculations for quantities like invariant mass or $\Delta R$.

Goal¶

In order to learn how to use these packages, we’ll try to do a very short and simple analysis to look for resonances in a dimuon events. The example ROOT file is from CMS open data from proton-proton collisions in 2012. Our first goal is just to get to the data within Uproot, so we need to open the file and navigate to the muon information.

Setup¶

Installation¶

These packages can be installed in the usual way if they haven't been already:

pip install uproot awkward vector

They are also available via conda install on the conda-forge channel.

Usage¶

We'll need NumPy and Matplotlib:

Uproot and Vector are normally imported by their name, and Awkward Array is usually imported as ak as an analogy to numpy as np:

Inspecting data¶

Files¶

Opening a file¶

Now open the ROOT file and assign it to a variable (which I’ve named file here):

If you inspect file, you can see that it’s a ReadOnlyDirectory object:

File contents¶

Just like any other kind of directory, you can list the contents (of the file). The name of each item in the file is called a key.

We can see that there is one key: “Events”. This doesn’t tell us what kind of object it refers to, though. ROOT files can contain many different types of objects, including subdirectories. The following function provides a way to inspect the types of each item:

The output contains pairs of the form name: type. Therefore the key Events refers to a TTree object. This is where all the data in this file is stored. The ;1 refers to the cycle number.

Accessing contents¶

Now we want to actually access the object inside the file. You can do this just as you would to get an item in an array:

This expression refers to the actual TTree object, which we will look at next.

Trees¶

Trees in ROOT are basically just tables of information. Trees are composed of branches, which are the columns of the table. The rows usually represent events (individual bunch crossings).

First we assign the tree to a variable (named tree here):

In order to find out what information is in the tree, we need to know what the branches (columns) are. The term key is used (again) here to refer to the names of the branches.

The above output is a list of the branch names. So we can see that for each event, we will have the number of muons in the event (nMuon) and the pT, eta, phi, mass, and charge of each muon. We can also see these keys with the types included using .show():

But how do we get the actual data from the table? There are several ways with Uproot, but the simplest is with the arrays() function:

You can see some numbers in there, which indeed are from the data in the tree.

Branches¶

Now we assign this object (which contains both the names and contents of the branches) to another variable (branches):

Next let’s just look at each branch individually. You can access a single branch from branches in a similar way to getting an item from a ROOT file object (array-like notation):

You can also access branches with attribute/property notation:

You can see the partial list of numbers in the output that represents the number of muons in each event. It’s abbreviated with an ellipsis (...) so that it doesn’t take up the whole page.

These Array objects are a special type provided by the Awkward Array package. The type=100000 * uint32 means that there are 100,000 entries and that each entry is a 32-bit unsigned integer. Each entry corresponds to one event.

Let’s look at another branch:

This is a jagged array because the number of entries is different for different events (because each event can have a different number of muons). Note that there are square brackets [] surrounding the list of entries for each event. The type='100000 * var * float32' means that there are 100,000 rows, each containing a variable number of 32-bit floating point numbers. This is basically an array of arrays (or a 2D array).

Events¶

If we want to focus on a particular event, we can index it just like a normal array:

From the above output, the first event has two muons, and the two numbers in the list are the muons’ $p_\mathrm{T}$. It’s not specified anywhere in the file, but the units are GeV. Let’s look at the third event:

It only has one muon.

What if we want to get all of the information about a single event? So far we’ve accessed data in branches by providing a branch name, but we can also just use an event index:

This is a Record object, which is another special type provided by Awkward Array. It functions basically the same way as a standard Python dictionary (dict). Unfortunately, most of the interesting information is still hidden in the above output to save space. A little trick we can use to force printing all the data is adding .tolist():

There we go. Now we can see the whole picture for an individual event.

.tolist() is a NumPy function that has been extended to Awkward Array objects. As the name suggests, it converts NumPy arrays to Python lists. In the case of trees, which have named branches, it actually converts to a dictionary of lists. It can be very useful when you want to understand exactly what’s in an Array or Record. Be careful when using it, though–trying to print out an entire branch or tree could cause Python to crash if it’s large enough. Therefore it’s best to only use tolist() on one or a few events at a time to be safe.

If you only want to see a description of the structure of the array and types within it (not the data itself), you can use .type:

Transforming data¶

Uproot is designed for columnar analysis, which means performing operations on entire columns (branches) at a time, rather than operating on every event individually.

Counting¶

The simplest task we need for analysis is counting (i.e., cutflow). To count the total number of events, we can use the Python built-in function len() on any of the following:

len(branches)
len(branches['nMuon'])
len(branches['Muon_pt']) # or any of the other branches...

which will output 100000. So there are 100,000 events.

Note that len will only tell you the length of the outermost dimension. The Muon_pt branch has more than 100,000 numbers in it. We can see this by flattening the array. np.ravel is a simple function to completely flatten an array:

Note that I've used a NumPy function here on an ak.Array. Most of the standard NumPy functions work fine on Awkward Arrays, although there are also equivalent Awkward functions like ak.ravel. Generally speaking, the default behavior between these NumPy and Awkward functions is the same, but the ak versions have optional parameters for funcationality that is specific to jagged arrays.

Awkward also has functions that wouldn't make sense in NumPy. For counting the number of entries in each event for a branch, there's ak.num:

which in this case is equal to already existing branches['nMuon']. In NumPy, there's no need for this function, since it would be the same number over and over, and you could just get it from np.ndarray.shape.

Selections¶

Selections from 1D arrays¶

To do more interesting counting or plotting, we want to be able to select only events or muons that pass some criteria. Just as you can compare two numbers, like 1 < 2 and get a True or False value, you can put arrays in comparisons as well. Let’s look at this case:

This is checking for equality between an array and the number 1. Of course these are different types (array vs. scalar), so the Python objects are certainly not identical, but you can see the return value is not just False. What happens is that each element in the array is compared to the scalar value, and the return value is a new array (of the same shape) filled with all these comparison results. So we can interpret the output as the result of testing whether each event has exactly one muon or not. The first two events do not, the third does, and so on.

This array of Boolean values is called a mask because we can use it to pick out only the elements in the array that satisfy some criteria (like having exactly one muon). This is very useful, and we will save it to a variable to save typing later:

Counting with selections¶

Now let’s say we want to know how many of these single-muon events there are. Note that len() won’t work because the length of the array is still 100,000. That is, there’s a value for every event (even if that value is False). We need the number of Trues in the array. There are multiple ways to do this, but my favorite is:

sum() adds all the array values together. True is interpreted as 1 and False is interpreted as 0, thus sum() is just the number of True entries. So there are 13,447 events with exactly one muon.

Applying a mask to an array¶

If we want to apply some selection to an array (i.e., cut out events or muons that don’t pass), we just act like the selection mask is an index. For example, if we want the pT of only those muons in events with exactly one muon (so they’re the only muon in that event):

We can check that we really are only looking at the events that pass the selection by looking the number of rows:

Yep, this matches the counting from single_muon_mask.sum() above.

Plotting with selections¶

We can also use masks to plot quantities after some selection. For example, let’s plot the muon $p_\mathrm{T}$ for only the single-muon events. Note that, just like for multidimensional NumPy arrays, Arrays must be flattened before histogramming.

Selections from a jagged array¶

Let’s look at a comparison for a jagged array, using the absolute value of muon eta:

Again, the mask array has the same dimensions as the original array. There’s one Boolean value for each muon, corresponding to whether its eta is less than 2 in absolute value.

We can do counting and plotting just as before:

This is the number of muons that pass the eta cut.

Plotting¶

Let’s plot both the original eta distribution and the one after the cut to verify its effect:

You can see the second plot just has both ends past 2 cut off, demonstrating that we’ve cut those muons out.

Operations on selections¶

We can invert selections with ~ (the NOT operator):

This new mask is False only for events with exactly one muon and True otherwise.

We can get the intersection of selections with & (the AND operator):

This mask is True only for muons with no other muons in their event and abs(eta) < 2.

Or we can get the union of selections with | (the OR operator):

This mask is True for muons which are the only muon in their event or which have abs(eta) < 2.

Comparing histograms¶

Now we can use these operations to compare distributions for different selections. Let’s look at the $p_\mathrm{T}$ of single-event muons split into two groups by whether or not abs(eta) < 2. All we have to do is provide a list of arrays as the first argument to hist rather than just one array. Note the square brackets around the two arrays:

Now we can clearly see there’s a significantly higher fraction of muons with abs(eta) >= 2 at lower pT compared to muons with abs(eta) < 2. This makes geometric sense, since muons at higher abs(eta) are traveling in a direction less perpendicular to the beam.

Analyzing data¶

Okay, we’re finally ready to look for resonances in dimuon events.

We need a mask that selects events with exactly two muons:

Vectors¶

Now we need to construct the four-momenta of each muon. Vector is a package that provides an interface to operate on 2D, 3D, and 4D vectors. We can get the four-momenta of all muons in the tree by using vector.zip and passing it the $p_\mathrm{T}$, eta, phi, and mass arrays:

We’ll go ahead and filter out events that don’t contain exactly two muons:

Then let’s take a look at it:

This is an array of 4D momenta. We can get the components back out with these properties:

What we want is the invariant mass of the two muons in each of these events. To do that, we need the sum of their four-vectors. First, we pick out the first muon in each event with 2D slice:

In the notation [:, 0], : means “include every row in the first dimension” (i.e., all events in the array). The comma separates the selection along the first dimension from the selection along the second dimension. The second dimension is the muons in each event, so we want the first, or the one at the 0 index. Then we do the same to get the second muon in each event, just changing the 0 index to 1:

$\Delta R$¶

Another useful feature of these four-vector arrays is being able to compute $\Delta R$ ($= \sqrt{{\Delta \eta}^2 + {\Delta \phi}^2}$):

In principle, we could use this to clean up our invariant mass distribution, but we’ll skip that for simplicity.

Dimuon mass¶

Adding the four-vectors of the first muon and the second muon for all events is really as easy as:

This is a 1D array of the four-vector sum for each event.

The last thing we need to do before we’re ready to plot the spectrum is to select only pairs with opposite charges:

We apply this selection to the four-vector sums to get the dimuon four-vectors:

Physics¶

Finally, we want to pick out the opposite-charge dimuon events and look at the distribution of invariant mass. This was our original physics goal in this toy analysis.

We create the necessary event mask:

Apply it to the four-vector sums:

And plot:

How many dimuon events were selected in total?

The plot should look like this, where the resonances have been labeled below:

Multiple files and large files¶

uproot.open() and tree.arrays() works fine for looking at one file at a time, but often trees are spread across several ROOT files. There are a few options to deal with this in Uproot.

The simplest is to join all the arrays:

uproot.concatenate({'*.root': 'Events'})

Another method is to iterate through the entries across all files:

for events in uproot.iterate({'*.root': 'Events'}, step_size=100):
    # do something with events

In each iteration of the loop, events is a 100-row chunk of the full tree. This avoids loading all entries at the same time. This is important if the files are too large to fit into memory.

There are also some other techniques to reduce the memory usage. If you are only interested in looking at a few events, you can do:

If you only need a subset of the branches from the tree, you can load each branch into an array individually, like:

or you can provide a list of branch names to arrays to get only those branches:

My preference is to use uproot.lazy():

Creating new objects and files¶

So far, we've seen how Uproot, Awkward, and Vector can be used to read and manipulate data from already existing ROOT files. In this last section we'll look at how these packages can also be used to construct new objects and write data to files.

Arrays¶

New Arrays can be easily be created like NumPy's np.array:

Vectors¶

Individual vectors can be created with vector.obj:

And we've already seen a way to make arrays of vectors:

Files¶

To write to a ROOT file, use uproot.create(), uproot.recreate(), or uproot.update(), which correspond to the same modes as in the option parameter of ROOT's TFile.Open(). For example, to create or replace new_file.root:

Writable file objects allow an assignment syntax for new elements, similar to a dictionary:

file['new_object_name'] = new_object

Trees¶

Trees can be specified by an Array or a dictionary of arrays:

And just to make sure we can read this back:

Histograms¶

Histograms from NumPy (or Hist) can be interpreted into ROOT's TH1. Let's save our dimuon mass histogram to this new ROOT file:

And again make sure we can read it back correctly:

Of course ROOT can read these newly created files as well, which makes it easy to drop in an Uproot-based replacement for any part of an analysis, even if later steps still expect ROOT files as input.

Links to more resources¶

• Uproot
  • GitHub repository
  • Documentation
• Awkward
  • GitHub repository
    • Contains several links to other presentations/tutorials
  • Documentation
• Vector
  • GitHub repository
  • Documentation