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