Pyroot 0 0 2_ T Tree As Matrix

This tutorial shows how a TTree can be quickly converted to a numpy array or a pandas.DataFrame.

Author: Stefan Wunsch
This notebook tutorial was automatically generated with ROOTBOOK-izer from the macro found in the ROOT repository on Saturday, June 06, 2020 at 02:44 AM.

In [1]:
import ROOT
from sys import exit

try:
    import numpy as np
except:
    print("Failed to import numpy.")
    exit()
Welcome to JupyROOT 6.23/01

Helper function to create an example tree

In [2]:
def make_example():
    root_file = ROOT.TFile("pyroot002_example.root", "RECREATE")
    tree = ROOT.TTree("tree", "tutorial")
    x = np.empty((1), dtype="float32")
    y = np.empty((1), dtype="float32")
    tree.Branch("x", x, "x/F")
    tree.Branch("y", y, "y/F")

    for i in range(4):
        x[0] = i
        y[0] = -i
        tree.Fill()
    root_file.Write()

    return (root_file, x, y), tree

The conversion of the TTree to a numpy array is implemented with multi- thread support.

In [3]:
ROOT.ROOT.EnableImplicitMT()

Create a ROOT file with a tree and the branches "x" and "y"

In [4]:
_, tree = make_example()

Print content of the tree by looping explicitly

In [5]:
print("Tree content:\n{}\n".format(
    np.asarray([[tree.x, tree.y] for event in tree])))
Tree content:
[[ 0.  0.]
 [ 1. -1.]
 [ 2. -2.]
 [ 3. -3.]]

Read-out full tree as numpy array

In [6]:
array = tree.AsMatrix()
print("Tree converted to a numpy array:\n{}\n".format(array))
Tree converted to a numpy array:
[[ 0.  0.]
 [ 1. -1.]
 [ 2. -2.]
 [ 3. -3.]]

Get numpy array and according labels of the columns

In [7]:
array, labels = tree.AsMatrix(return_labels=True)
print("Return numpy array and labels:\n{}\n{}\n".format(labels, array))
Return numpy array and labels:
['x', 'y']
[[ 0.  0.]
 [ 1. -1.]
 [ 2. -2.]
 [ 3. -3.]]

Apply numpy methods on the data

In [8]:
print("Mean of the columns retrieved with a numpy method: {}\n".format(
    np.mean(array, axis=0)))
Mean of the columns retrieved with a numpy method: [ 1.5 -1.5]

Read only specific branches

In [9]:
array = tree.AsMatrix(columns=["x"])
print("Only the content of the branch 'x':\n{}\n".format(np.squeeze(array)))

array = tree.AsMatrix(exclude=["x"])
print("Read all branches except 'x':\n{}\n".format(np.squeeze(array)))
Only the content of the branch 'x':
[0. 1. 2. 3.]

Read all branches except 'x':
[ 0. -1. -2. -3.]

Get an array with a specific data-type

In [10]:
array = tree.AsMatrix(dtype="int")
print("Return numpy array with data-type 'int':\n{}\n".format(array))
Return numpy array with data-type 'int':
[[ 0  0]
 [ 1 -1]
 [ 2 -2]
 [ 3 -3]]

Convert the tree to a pandas.DataFrame

In [11]:
try:
    import pandas
except:
    print("Failed to import pandas.")
    exit()

data, columns = tree.AsMatrix(return_labels=True)
df = pandas.DataFrame(data=data, columns=columns)
print("Tree converted to a pandas.DataFrame:\n{}".format(df))
Tree converted to a pandas.DataFrame:
     x    y
0  0.0  0.0
1  1.0 -1.0
2  2.0 -2.0
3  3.0 -3.0