ROOT   Reference Guide
df002_dataModel.py
Go to the documentation of this file.
1## \file
2## \ingroup tutorial_dataframe
3## \notebook -draw
4## Show how to work with non-flat data models, e.g. vectors of tracks.
5##
6## This tutorial shows the possibility to use data models which are more
7## complex than flat ntuples with RDataFrame
8##
9## \macro_code
10## \macro_image
11##
12## \date May 2017
13## \author Danilo Piparo (CERN)
14
15import ROOT
16
17# A simple helper function to fill a test tree: this makes the example stand-alone.
18fill_tree_code = '''
19using FourVector = ROOT::Math::XYZTVector;
20using FourVectorVec = std::vector<FourVector>;
21using CylFourVector = ROOT::Math::RhoEtaPhiVector;
22
23// A simple helper function to fill a test tree: this makes the example
24// stand-alone.
25void fill_tree(const char *filename, const char *treeName)
26{
27 const double M = 0.13957; // set pi+ mass
28 TRandom3 R(1);
29
30 auto genTracks = [&](){
31 FourVectorVec tracks;
32 const auto nPart = R.Poisson(15);
33 tracks.reserve(nPart);
34 for (int j = 0; j < nPart; ++j) {
35 const auto px = R.Gaus(0, 10);
36 const auto py = R.Gaus(0, 10);
37 const auto pt = sqrt(px * px + py * py);
38 const auto eta = R.Uniform(-3, 3);
39 const auto phi = R.Uniform(0.0, 2 * TMath::Pi());
40 CylFourVector vcyl(pt, eta, phi);
41 // set energy
42 auto E = sqrt(vcyl.R() * vcyl.R() + M * M);
43 // fill track vector
44 tracks.emplace_back(vcyl.X(), vcyl.Y(), vcyl.Z(), E);
45 }
46 return tracks;
47 };
48
49 ROOT::RDataFrame d(64);
50 d.Define("tracks", genTracks).Snapshot<FourVectorVec>(treeName, filename, {"tracks"});
51}
52'''
53
54# We prepare an input tree to run on
55fileName = "df002_dataModel_py.root"
56treeName = "myTree"
57ROOT.gInterpreter.Declare(fill_tree_code)
58ROOT.fill_tree(fileName, treeName)
59
60# We read the tree from the file and create a RDataFrame, a class that
61# allows us to interact with the data contained in the tree.
62d = ROOT.RDataFrame(treeName, fileName)
63
64# Operating on branches which are collection of objects
65# Here we deal with the simplest of the cuts: we decide to accept the event
66# only if the number of tracks is greater than 8.
67n_cut = 'tracks.size() > 8'
68nentries = d.Filter(n_cut).Count();
69
70print("%s events passed all filters" % nentries.GetValue())
71
72# Another possibility consists in creating a new column containing the
73# quantity we are interested in.
74# In this example, we will cut on the number of tracks and plot their
75# transverse momentum.
76
77getPt_code ='''
78using namespace ROOT::VecOps;
79RVec<double> getPt(const RVec<FourVector> &tracks)
80{
81 auto pt = [](const FourVector &v) { return v.pt(); };
82 return Map(tracks, pt);
83}
84'''
85ROOT.gInterpreter.Declare(getPt_code)
86
87getPtWeights_code ='''
88using namespace ROOT::VecOps;
89RVec<double> getPtWeights(const RVec<FourVector> &tracks)
90{
91 auto ptWeight = [](const FourVector &v) { return 1. / v.Pt(); };
92 return Map(tracks, ptWeight);
93};
94'''
95ROOT.gInterpreter.Declare(getPtWeights_code)
96
97augmented_d = d.Define('tracks_n', '(int)tracks.size()') \
98 .Filter('tracks_n > 2') \
99 .Define('tracks_pts', 'getPt( tracks )') \
100 .Define("tracks_pts_weights", 'getPtWeights( tracks )' )
101
102# The histogram is initialised with a tuple containing the parameters of the
103# histogram
104trN = augmented_d.Histo1D(("", "", 40, -.5, 39.5), "tracks_n")
105trPts = augmented_d.Histo1D("tracks_pts")
106trWPts = augmented_d.Histo1D("tracks_pts", "tracks_pts_weights")
107
108c1 = ROOT.TCanvas()
109trN.Draw()
110c1.SaveAs("df002_trN.png")
111
112c2 = ROOT.TCanvas()
113trPts.Draw()
114c2.SaveAs("df002_trPts.png")
115
116c3 = ROOT.TCanvas()
117trWPts.Draw()
118c2.SaveAs("df002_trWPts.png")
119
120print("Saved figures to df002_*.png")
#define d(i)
Definition: RSha256.hxx:102
double sqrt(double)
double Pi()
Mathematical constants.
Definition: Math.h:88
auto Map(Args &&... args) -> decltype(ROOT::Detail::VecOps::MapFromTuple(std::forward_as_tuple(args...), std::make_index_sequence< sizeof...(args) - 1 >()))
Create new collection applying a callable to the elements of the input collection.
Definition: RVec.hxx:942
RVec< T > Filter(const RVec< T > &v, F &&f)
Create a new collection with the elements passing the filter expressed by the predicate.
Definition: RVec.hxx:971