Logo ROOT  
Reference Guide
pyroot002_TTreeAsMatrix.py
Go to the documentation of this file.
1## \file
2## \ingroup tutorial_pyroot
3## \notebook -nodraw
4## This tutorial shows how a TTree can be quickly converted to a numpy array or
5## a pandas.DataFrame.
6##
7## \macro_code
8## \macro_output
9##
10## \date April 2018
11## \author Stefan Wunsch
12
13import ROOT
14from sys import exit
15
16try:
17 import numpy as np
18except:
19 print("Failed to import numpy.")
20 exit()
21
22
23# Helper function to create an example tree
24def make_example():
25 root_file = ROOT.TFile("pyroot002_example.root", "RECREATE")
26 tree = ROOT.TTree("tree", "tutorial")
27 x = np.empty((1), dtype="float32")
28 y = np.empty((1), dtype="float32")
29 tree.Branch("x", x, "x/F")
30 tree.Branch("y", y, "y/F")
31
32 for i in range(4):
33 x[0] = i
34 y[0] = -i
35 tree.Fill()
36 root_file.Write()
37
38 return (root_file, x, y), tree
39
40
41# The conversion of the TTree to a numpy array is implemented with multi-
42# thread support.
43ROOT.ROOT.EnableImplicitMT()
44
45# Create a ROOT file with a tree and the branches "x" and "y"
46_, tree = make_example()
47
48# Print content of the tree by looping explicitly
49print("Tree content:\n{}\n".format(
50 np.asarray([[tree.x, tree.y] for event in tree])))
51
52# Read-out full tree as numpy array
53array = tree.AsMatrix()
54print("Tree converted to a numpy array:\n{}\n".format(array))
55
56# Get numpy array and according labels of the columns
57array, labels = tree.AsMatrix(return_labels=True)
58print("Return numpy array and labels:\n{}\n{}\n".format(labels, array))
59
60# Apply numpy methods on the data
61print("Mean of the columns retrieved with a numpy method: {}\n".format(
62 np.mean(array, axis=0)))
63
64# Read only specific branches
65array = tree.AsMatrix(columns=["x"])
66print("Only the content of the branch 'x':\n{}\n".format(np.squeeze(array)))
67
68array = tree.AsMatrix(exclude=["x"])
69print("Read all branches except 'x':\n{}\n".format(np.squeeze(array)))
70
71# Get an array with a specific data-type
72array = tree.AsMatrix(dtype="int")
73print("Return numpy array with data-type 'int':\n{}\n".format(array))
74
75## Convert the tree to a pandas.DataFrame
76try:
77 import pandas
78except:
79 print("Failed to import pandas.")
80 exit()
81
82data, columns = tree.AsMatrix(return_labels=True)
83df = pandas.DataFrame(data=data, columns=columns)
84print("Tree converted to a pandas.DataFrame:\n{}".format(df))