I think it would be nice if the adding of a new branch to a tree could be simplified. This comes up a lot, because most analysis start with a selection script that also calculates some new branches based on the already existing branches. Since one only wants to use that new processed root file afterwards using friends is not really an option (as opposed to the case when calculating weights for example).
In order to do this, I had already posted a script that could do this for multiple branches (see here), but I found out that doing one branch at a time is usually not slower, but a lot easier. The following function does just that (see example at the end).
template <typename T>
void AddBranch(const char *name, const char *expression, TTree *tree)
{
T x;
auto branch = tree->Branch(name, &x);
auto formula = new TTreeFormula(name, expression, tree);
// Activate only needed branches
tree->SetBranchStatus("*", kFALSE);
tree->SetBranchStatus(name, kTRUE);
for (Int_t i = 0; i < formula->GetNcodes(); i++) {
tree->SetBranchStatus(formula->GetLeaf(i)->GetName(), kTRUE);
}
// Calculate content of new branch
Long64_t nentries = tree->GetEntries();
for (Long64_t i = 0; i < nentries; i++) {
if (tree->GetEntry(i)) {
x = formula->EvalInstance();
branch->Fill();
}
}
// re-activate all branches and reset their addresses
tree->SetBranchStatus("*", kTRUE);
// Maybe this is overkill, because it resets all?
tree->ResetBranchAddresses();
}
Example usage:
$ root $ROOTSYS/tutorials/hsimple.root
root [0]
Attaching file /cern/root/tutorials/hsimple.root as _file0...
(TFile *) 0x55c1a6eec110
root [1] auto tree = (TTree*) _file0->Get("ntuple")
(TTree *) 0x55c1a6d636e0
root [3] auto fout = new TFile("out.root", "recreate")
(TFile *) 0x55c1a7a35d40
root [5] auto tout = tree->CopyTree("random > 0.5")
(TTree *) 0x55c1a7a2de40
root [8] #include "/path/to/utils.h"
root [9] AddBranch<Double_t>("P", "px**2+py**2+pz**2", tout)
root [10] tout->Write()
(int) 964
root [11] fout->Write()
(int) 964
root [12] tout->Draw("P")
Info in <TCanvas::MakeDefCanvas>: created default TCanvas with name c1
root [13] fout->Close()
The example doesn’t compile with ROOT 6.09.02, though, I guess I need the master branch for that?
$ root tdf007_snapshot.cpp++
root [0]
Processing tdf007_snapshot.cpp++...
Info in <TUnixSystem::ACLiC>: creating shared library ./tdf007_snapshot_cpp.so
In file included from input_line_11:9:
././tdf007_snapshot.cpp:40:20: error: no member named 'Define' in 'ROOT::Experimental::TDataFrameInterface<ROOT::Detail::TDataFrameFilterBase>'
auto d2 = d_cut.Define("b1_square", "b1 * b1")
~~~~~ ^
././tdf007_snapshot.cpp:80:39: error: expected '(' for function-style cast or type construction
auto snapshot_tdf = d2.Snapshot<int>(treeName, outFileName, {"b1_square"});
~~~^
././tdf007_snapshot.cpp:80:41: warning: expression result unused [-Wunused-value]
auto snapshot_tdf = d2.Snapshot<int>(treeName, outFileName, {"b1_square"});
^~~~~~~~
././tdf007_snapshot.cpp:80:62: error: expected ')'
auto snapshot_tdf = d2.Snapshot<int>(treeName, outFileName, {"b1_square"});
^
././tdf007_snapshot.cpp:80:40: note: to match this '('
auto snapshot_tdf = d2.Snapshot<int>(treeName, outFileName, {"b1_square"});
^
In file included from input_line_11:9:
In file included from ././tdf007_snapshot.cpp:5:
In file included from /cern/root-6.09.02/include/ROOT/TDataFrame.hxx:21:
In file included from /cern/root-6.09.02/include/ROOT/TDFOperations.hxx:14:
/cern/root-6.09.02/include/ROOT/TDFTraitsUtils.hxx:31:54: error: type 'const char *' cannot be used prior to '::' because it has no members
using Args_t = typename TFunctionTraits<decltype(&T::operator())>::Args_t;
^
/cern/root-6.09.02/include/ROOT/TDataFrame.hxx:421:52: note: in instantiation of template class 'ROOT::Internal::TDFTraitsUtils::TFunctionTraits<const char *>' requested here
auto nArgs = ROOT::Internal::TDFTraitsUtils::TFunctionTraits<F>::Args_t::fgSize;
^
././tdf007_snapshot.cpp:37:19: note: in instantiation of function template specialization
'ROOT::Experimental::TDataFrameInterface<ROOT::Detail::TDataFrameImpl>::Filter<const char *>' requested here
auto d_cut = d.Filter("b1 % 2 == 0");
^
In file included from input_line_11:9:
In file included from ././tdf007_snapshot.cpp:5:
In file included from /cern/root-6.09.02/include/ROOT/TDataFrame.hxx:21:
In file included from /cern/root-6.09.02/include/ROOT/TDFOperations.hxx:14:
/cern/root-6.09.02/include/ROOT/TDFTraitsUtils.hxx:32:61: error: type 'const char *' cannot be used prior to '::' because it has no members
using ArgsNoDecay_t = typename TFunctionTraits<decltype(&T::operator())>::ArgsNoDecay_t;
^
/cern/root-6.09.02/include/ROOT/TDFTraitsUtils.hxx:33:53: error: type 'const char *' cannot be used prior to '::' because it has no members
using Ret_t = typename TFunctionTraits<decltype(&T::operator())>::Ret_t;
^
In file included from input_line_11:9:
In file included from ././tdf007_snapshot.cpp:5:
/cern/root-6.09.02/include/ROOT/TDataFrame.hxx:215:4: error: static_assert failed "filter functions must return a bool"
static_assert(std::is_same<FilterRet_t, bool>::value, "filter functions must return a bool");
^ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/cern/root-6.09.02/include/ROOT/TDataFrame.hxx:418:23: note: in instantiation of function template specialization 'ROOT::Internal::CheckFilter<const char *>' requested here
ROOT::Internal::CheckFilter(f);
^
././tdf007_snapshot.cpp:37:19: note: in instantiation of function template specialization
'ROOT::Experimental::TDataFrameInterface<ROOT::Detail::TDataFrameImpl>::Filter<const char *>' requested here
auto d_cut = d.Filter("b1 % 2 == 0");
^
Error in <ACLiC>: Dictionary generation failed!
and way more complete than what I tried.