import ROOT
fill_tree_code ='''
void fill_tree(const char *filename, const char *treeName)
{
TFile f(filename, "RECREATE");
TTree t(treeName, treeName);
int b1;
float b2;
t.Branch("b1", &b1);
t.Branch("b2", &b2);
for (int i = 0; i < 10000; ++i) {
b1 = i;
b2 = i * i;
t.Fill();
}
t.Write();
f.Close();
return;
}
'''
fileName = "tdf007_snapshot_py.root"
outFileName = "tdf007_snapshot_output_py.root"
outFileNameAllColumns = "tdf007_snapshot_output_allColumns_py.root"
treeName = "myTree"
ROOT.gInterpreter.Declare(fill_tree_code)
ROOT.fill_tree(fileName, treeName)
TDF = ROOT.ROOT.Experimental.TDataFrame
d = TDF(treeName, fileName)
d_cut = d.Filter("b1 % 2 == 0")
getVector_code ='''
std::vector<float> getVector (float b2)
{
std::vector<float> v;
for (int i = 0; i < 3; i++) v.push_back(b2*i);
return v;
}
'''
ROOT.gInterpreter.Declare(getVector_code)
d2 = d_cut.Define("b1_square", "b1 * b1") \
.Define("b2_vector", "getVector( b2 )")
branchList = ROOT.vector('string')()
for branchName in ["b1", "b1_square", "b2_vector"]:
branchList.push_back(branchName)
d2.Snapshot(treeName, outFileName, branchList)
f1 = ROOT.TFile(outFileName)
t = f1.myTree
print("These are the columns b1, b1_square and b2_vector:")
for branch in t.GetListOfBranches():
print("Branch: %s" %branch.GetName())
f1.Close()
d2.Snapshot(treeName, outFileNameAllColumns)
f2 = ROOT.TFile(outFileNameAllColumns)
t = f2.myTree
print("These are all the columns available to this tdf:")
for branch in t.GetListOfBranches():
print("Branch: %s" %branch.GetName())
f2.Close()
branchList.clear()
branchList.push_back("b1_square")
snapshot_tdf = d2.Snapshot(treeName, outFileName, branchList);
h = snapshot_tdf.Histo1D("b1_square")
c = ROOT.TCanvas()
h.Draw()