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 
13 import ROOT
14 from sys import exit
15 
16 try:
17  import numpy as np
18 except:
19  print("Failed to import numpy.")
20  exit()
21 
22 
23 # Helper function to create an example tree
24 def 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.
43 ROOT.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
49 print("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
53 array = tree.AsMatrix()
54 print("Tree converted to a numpy array:\n{}\n".format(array))
55 
56 # Get numpy array and according labels of the columns
57 array, labels = tree.AsMatrix(return_labels=True)
58 print("Return numpy array and labels:\n{}\n{}\n".format(labels, array))
59 
60 # Apply numpy methods on the data
61 print("Mean of the columns retrieved with a numpy method: {}\n".format(
62  np.mean(array, axis=0)))
63 
64 # Read only specific branches
65 array = tree.AsMatrix(columns=["x"])
66 print("Only the content of the branch 'x':\n{}\n".format(np.squeeze(array)))
67 
68 array = tree.AsMatrix(exclude=["x"])
69 print("Read all branches except 'x':\n{}\n".format(np.squeeze(array)))
70 
71 # Get an array with a specific data-type
72 array = tree.AsMatrix(dtype="int")
73 print("Return numpy array with data-type 'int':\n{}\n".format(array))
74 
75 ## Convert the tree to a pandas.DataFrame
76 try:
77  import pandas
78 except:
79  print("Failed to import pandas.")
80  exit()
81 
82 data, columns = tree.AsMatrix(return_labels=True)
83 df = pandas.DataFrame(data=data, columns=columns)
84 print("Tree converted to a pandas.DataFrame:\n{}".format(df))