Skip to content

Instantly share code, notes, and snippets.

@JavierCVilla
Last active October 25, 2018 07:55
Show Gist options
  • Save JavierCVilla/1836f2f3e5188cd5524c9113f6dff244 to your computer and use it in GitHub Desktop.
Save JavierCVilla/1836f2f3e5188cd5524c9113f6dff244 to your computer and use it in GitHub Desktop.
Script to create multiple trees split on different files.
import ROOT
import sys
# A simple helper function to fill a test tree: this makes the example stand-alone.
def fill_tree(treeName, fileName, b1, b2):
opts = ROOT.ROOT.RDF.RSnapshotOptions();
opts.fMode = "UPDATE";
branchList = ROOT.vector('string')()
for branchName in [b1, b2]:
branchList.push_back(branchName)
tdf = ROOT.ROOT.RDataFrame(200)
tdf.Define(b1, "(double) tdfentry_")\
.Define(b2, "(int) tdfentry_ * tdfentry_")\
.Snapshot(treeName, fileName, branchList, opts)
# We prepare some trees split on different files
for f in range(20):
filename = "sample{}.root".format(f)
fill_tree("processInfoTree", filename, "luminosity", "cross_section")
fill_tree("weightsTree", filename, "weight", "weight1")
fill_tree("randomTree", filename, "x", "x2")
fill_tree("anotherTree", filename, "y", "y2")
{
using namespace ROOT;
auto filenames = std::vector<std::string>();
for(int i = 0; i<20 ; i++){
std::string s = "sample" + std::to_string(i) + ".root";
filenames.push_back(s);
}
// process info tree
TChain *info_chain = new TChain("processInfoTree");
// weight tree
TChain *weight_chain = new TChain("weightsTree");
// weight 2 tree
TChain *weight2_chain = new TChain("randomTree");
// weight 3 tree
TChain *weight3_chain = new TChain("anotherTree");
for(auto f: filenames){
cout << f << endl;
info_chain->Add(f.c_str());
weight_chain->Add(f.c_str());
weight2_chain->Add(f.c_str());
weight3_chain->Add(f.c_str());
}
info_chain->AddFriend(weight_chain, "weights");
info_chain->AddFriend(weight2_chain, "weights");
info_chain->AddFriend(weight3_chain, "weights");
RDataFrame* rdf = new RDataFrame(*info_chain);
cout << "Branches: " << endl;
for(auto c : rdf->GetColumnNames())
cout << c << endl;
cout << *rdf->Filter("cross_section > 0").Count() << endl;
cout << *rdf->Filter("weight1 > 0.5").Count() << endl; // Fails here
delete rdf;
delete info_chain;
delete weight_chain;
}
@JavierCVilla
Copy link
Author

JavierCVilla commented Oct 25, 2018

Using different aliases for the friend trees I did not manage to reproduce the problem:

info_chain->AddFriend(weight_chain, "weights");
info_chain->AddFriend(weight2_chain, "weights2");
info_chain->AddFriend(weight3_chain, "weights3");

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment