Skip to content

Commit

Permalink
Unified tree writing and adding of formula branches
Browse files Browse the repository at this point in the history
  • Loading branch information
@VukanJ
VukanJ committed Jan 20, 2019
1 parent 3a87cc0 commit a5630d8
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 83 deletions.
124 changes: 63 additions & 61 deletions Ranger.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,13 +115,8 @@ void Ranger::BPVselection(const std::string& tree_in,

void Ranger::addFormula(const std::string& name, std::string formula)
{
if (tree_jobs.empty()) {
std::cerr << "Error: Need a previous tree job for adding a formula branch! Skipping\n";
return;
}
tree_jobs.push_back({
{{"tree_out", tree_jobs.back()["tree_out"]},
{"formula", formula},
{{"formula", formula},
{"branch_name", name}},
Action::add_formula});
}
Expand All @@ -141,13 +136,13 @@ void Ranger::Run(TString output_filename)
}
closeFile(outFile.get());
outFile = FilePtr(TFile::Open(output_filename, "RECREATE"));

for (auto& tree_job : tree_jobs) {
switch (tree_job.action) {
case Action::copytree: SimpleCopy(tree_job); break;
case Action::flatten_tree: flattenTree(tree_job); break;
case Action::bpv_selection: BestPVSelection(tree_job); break;
case Action::add_formula: addFormulaBranch(tree_job); break;
case Action::add_formula: formula_buffer.push_back(std::make_pair(tree_job["branch_name"],
tree_job["formula"])); break;
default: break;
}
clearBuffers();
Expand All @@ -163,11 +158,46 @@ void Ranger::reset()
clearBuffers();
}

void Ranger::finalizeTreeWriting(const TreeJob& tree_job, TTree* tree, bool tree_prepared)
{
// Add formula branches, apply cuts, write to file
if (!formula_buffer.empty()) {
TTree* temp_tree = tree->CloneTree();
for (const auto& formula : formula_buffer){
addFormulaBranch(temp_tree, formula.first, formula.second);
}
formula_buffer.clear();
if (tree_job["cut"] == ""){
temp_tree->Write("", TObject::kOverwrite);
}
else {
TTree* output_tree = temp_tree->CopyTree(tree_job("cut"));
output_tree->Write("", TObject::kOverwrite);
delete output_tree;
}
delete temp_tree;
}
else if (tree_job["cut"] != ""){
TTree* output_tree = tree->CopyTree(tree_job("cut"));
output_tree->Write("", TObject::kOverwrite);
delete output_tree;
}
else {
if(!tree_prepared){
TTree* output_tree = tree->CloneTree();
output_tree->Write("", TObject::kOverwrite);
delete output_tree;
}
else{
tree->Write("", TObject::kOverwrite);
}
}
}

void Ranger::SimpleCopy(const TreeJob& tree_job)
{
// Copy tree with cut selection and branch selection using built-in methods
TTree* input_tree = static_cast<TTree*>(inFile->Get(tree_job("tree_in")));
TTree* output_tree = nullptr;
TTree* input_tree = static_cast<TTree*>(inFile->Get(tree_job("tree_in")));

if (!tree_job["branch_selection"].empty()) {
input_tree->SetBranchStatus("*", 0);
Expand All @@ -180,17 +210,9 @@ void Ranger::SimpleCopy(const TreeJob& tree_job)
else {
input_tree->SetBranchStatus("*", 1);
}

if (tree_job["cut"] == "") {
output_tree = static_cast<TTree*>(input_tree->CloneTree());
}
else {
output_tree = static_cast<TTree*>(input_tree->CopyTree(tree_job("cut")));
}
output_tree->SetTitle(tree_job("tree_out"));

output_tree->Write("", TObject::kOverwrite); // Disable Autosave backups
std::cout << "Done copying\n";
input_tree->SetTitle(tree_job("tree_out"));
input_tree->SetName(tree_job("tree_out"));
finalizeTreeWriting(tree_job, input_tree, false);
}

void Ranger::flattenTree(const TreeJob& tree_job)
Expand Down Expand Up @@ -223,31 +245,24 @@ void Ranger::flattenTree(const TreeJob& tree_job)
for (int event = 0; event < n_entries; ++event) {
input_tree->GetEntry(event);
array_elem_it = 0;
std::cout << "EVENT " << event << ' ' << max_array_length << '\n';
//std::cout << "EVENT " << event << ' ' << max_array_length << '\n';
output_tree.Fill();
for (array_elem_it = 1; array_elem_it < max_array_length; ++array_elem_it) {
// go to next array element
incrementBuffer<Char_t>(array_elem_it);
incrementBuffer<UChar_t>(array_elem_it);
incrementBuffer<Short_t>(array_elem_it);
incrementBuffer<UShort_t>(array_elem_it);
incrementBuffer<Int_t>(array_elem_it);
incrementBuffer<UInt_t>(array_elem_it);
incrementBuffer<Double_t>(array_elem_it);
incrementBuffer<Float_t>(array_elem_it);
incrementBuffer<Long64_t>(array_elem_it);
incrementBuffer< Char_t>(array_elem_it);
incrementBuffer< UChar_t>(array_elem_it);
incrementBuffer< Short_t>(array_elem_it);
incrementBuffer< UShort_t>(array_elem_it);
incrementBuffer< Int_t>(array_elem_it);
incrementBuffer< UInt_t>(array_elem_it);
incrementBuffer< Double_t>(array_elem_it);
incrementBuffer< Float_t>(array_elem_it);
incrementBuffer< Long64_t>(array_elem_it);
incrementBuffer<ULong64_t>(array_elem_it);
output_tree.Fill();
}
}
if (!tree_job["cut"].empty()) {
//Create intermediate tree for copying with selection
auto* output_tree_selected = static_cast<TTree*>(output_tree.CopyTree(tree_job("cut")));
output_tree_selected->Write("", TObject::kOverwrite);
}
else {
output_tree.Write("", TObject::kOverwrite);
}
finalizeTreeWriting(tree_job, &output_tree);
}

void Ranger::BestPVSelection(const TreeJob& tree_job)
Expand Down Expand Up @@ -275,14 +290,7 @@ void Ranger::BestPVSelection(const TreeJob& tree_job)
input_tree->GetEntry(event);
output_tree.Fill();
}
if (!tree_job["cut"].empty()) {
//Create intermediate tree for copying with selection
TTree* output_tree_selected = static_cast<TTree*>(output_tree.CopyTree(tree_job("cut")));
output_tree_selected->Write("", TObject::kOverwrite);
}
else {
output_tree.Write("", TObject::kOverwrite);
}
finalizeTreeWriting(tree_job, &output_tree);
}

TLeaf* Ranger::analyzeLeaves_FillLeafBuffers(TTree* input_tree, TTree* output_tree,
Expand Down Expand Up @@ -417,16 +425,11 @@ void Ranger::getListOfBranchesBySelection(std::vector<TLeaf*>& selected, TTree*
}
}

void Ranger::addFormulaBranch(const TreeJob& tree_job)
void Ranger::addFormulaBranch(TTree* output_tree, const std::string& name, std::string formula)
{
TTree* output_tree_existing = static_cast<TTree*>(outFile->Get(tree_job("tree_out")));
std::string formula = tree_job["formula"]; // Name does not really fit here ...

std::regex var_search(R"(\#[\w_][\w\d_]*)"); // Matches variables

// Extract variable strings from formula string
std::set<std::string> variables;

// Extract variables from formula string
std::regex var_search(R"(\#[\w_][\w\d_]*)"); // Matches variables
std::sregex_iterator iter(formula.begin(), formula.end(), var_search);
std::sregex_iterator end;
for (; iter != end; ++iter) {
Expand All @@ -436,25 +439,24 @@ void Ranger::addFormulaBranch(const TreeJob& tree_job)
std::vector<Double_t> buffer(variables.size());

Double_t result;
TBranch* formula_branch = output_tree_existing->Branch(tree_job("branch_name"), &result);
TBranch* formula_branch = output_tree->Branch(TString(name), &result);

int idx = 0;
for (std::string var : variables) {
formula = std::regex_replace(formula, std::regex(var), std::string("[") + std::to_string(idx) + "]");
var.erase(var.begin()); // Remove '#'
output_tree_existing->SetBranchStatus(TString(var), 1);
output_tree_existing->SetBranchAddress(TString(var), &buffer[idx]);
output_tree->SetBranchStatus(TString(var), 1);
output_tree->SetBranchAddress(TString(var), &buffer[idx]);
++idx;
}

TFormula tformula("F", TString(formula));

int n_entries = output_tree_existing->GetEntriesFast();
int n_entries = output_tree->GetEntriesFast();
for (int event = 0; event < n_entries; ++event) {
output_tree_existing->GetEntry(event);
output_tree->GetEntry(event);
result = tformula.EvalPar(nullptr, &buffer[0]);

formula_branch->Fill();
}
output_tree_existing->Write("", TObject::kOverwrite);
}
25 changes: 15 additions & 10 deletions Ranger.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ class Ranger {
// Utility methods
void closeFile(TFile*);
void clearBuffers();
void finalizeTreeWriting(const TreeJob&, TTree*, bool tree_prepared=true);
TLeaf* analyzeLeaves_FillLeafBuffers(TTree* input_tree, TTree* output_tree,
std::vector<TLeaf*>& all_leaves,
std::vector<TLeaf*>& bpv_leaves);
Expand All @@ -113,7 +114,9 @@ class Ranger {
void SimpleCopy(const TreeJob&);
void flattenTree(const TreeJob&);
void BestPVSelection(const TreeJob&);
void addFormulaBranch(const TreeJob&);
void addFormulaBranch(TTree* output_tree,
const std::string& name,
std::string formula);

std::vector<TreeJob> tree_jobs;

Expand All @@ -133,22 +136,24 @@ class Ranger {
Buffer<Long64_t> leaf_buffers_L;
Buffer<ULong64_t> leaf_buffers_l;

std::vector<std::pair<std::string, std::string>> formula_buffer;

ClassDef(Ranger,1)
};

template<typename L>
Buffer<L>* Ranger::getBuffer()
{
// Return buffer corresponding to datatype
if constexpr (std::is_same<L, Char_t>::value) return &leaf_buffers_B;
if constexpr (std::is_same<L, UChar_t>::value) return &leaf_buffers_b;
if constexpr (std::is_same<L, Short_t>::value) return &leaf_buffers_S;
if constexpr (std::is_same<L, UShort_t>::value) return &leaf_buffers_s;
if constexpr (std::is_same<L, Int_t>::value) return &leaf_buffers_I;
if constexpr (std::is_same<L, UInt_t>::value) return &leaf_buffers_i;
if constexpr (std::is_same<L, Float_t>::value) return &leaf_buffers_F;
if constexpr (std::is_same<L, Double_t>::value) return &leaf_buffers_D;
if constexpr (std::is_same<L, Long64_t>::value) return &leaf_buffers_L;
if constexpr (std::is_same<L, Char_t>::value) return &leaf_buffers_B;
if constexpr (std::is_same<L, UChar_t>::value) return &leaf_buffers_b;
if constexpr (std::is_same<L, Short_t>::value) return &leaf_buffers_S;
if constexpr (std::is_same<L, UShort_t>::value) return &leaf_buffers_s;
if constexpr (std::is_same<L, Int_t>::value) return &leaf_buffers_I;
if constexpr (std::is_same<L, UInt_t>::value) return &leaf_buffers_i;
if constexpr (std::is_same<L, Float_t>::value) return &leaf_buffers_F;
if constexpr (std::is_same<L, Double_t>::value) return &leaf_buffers_D;
if constexpr (std::is_same<L, Long64_t>::value) return &leaf_buffers_L;
if constexpr (std::is_same<L, ULong64_t>::value) return &leaf_buffers_l;
return nullptr;
}
Expand Down
16 changes: 5 additions & 11 deletions dev.py
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,9 @@
from root_ranger import Ranger
import time

ranger = Ranger("../test/0/DTT.root")
start = time.time()
#ranger.bpv_selection("inclusive_Jpsi/DecayTree", "B0_Fit*", dest="DecayTree")
ranger.flatten_tree("inclusive_Jpsi/DecayTree", "B0_Fit*", branches="B0*", dest="DecayTree")
ranger.run("Test_flat.root")
end = time.time()
ranger.add_formula("F", "TMath::Sqrt(#B0_M)")
ranger.copy_tree("inclusive_Jpsi/DecayTree", branches="B0_PX", dest="MyTree")
#ranger.bpv_selection("inclusive_Jpsi/DecayTree", "B0_Fit*", dest="MyTree")
ranger.run("DTT_flat.root")


print("Took, ", end-start, " seconds")
#ranger2 = Ranger("../test/0/DTT.root")
#
#ranger2.run("Test_flat.root")
# Invalid input tree names crashes program in a non-trivial way!
2 changes: 1 addition & 1 deletion root_ranger.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def bpv_selection(self, treename, bpv_branches, branches="*", cut="", dest=""):
self.__extend_selection(cut),
dest)

def add_formula(formula_name, formula):
def add_formula(self, formula_name, formula):
# For example add_formula("B0_PT", "TMath::Sqrt(#B0_X**2+#B0_Y**2)")
self.__ranger.addFormula(formula_name, formula)

Expand Down

0 comments on commit a5630d8

Please sign in to comment.