diff --git a/rapido/src/arbusto.h b/rapido/src/arbusto.h index 86796be..0e10d91 100644 --- a/rapido/src/arbusto.h +++ b/rapido/src/arbusto.h @@ -12,8 +12,10 @@ class Arbusto : public Arbol { protected: + /** Invert how the branches are kept */ + bool remove_branches; /** Names of branches in original ROOT TTree to keep */ - std::vector keep_branch_names; + std::vector branch_names; /** Pointer to current TTree to skim */ TTree* orig_ttree; public: @@ -27,9 +29,10 @@ class Arbusto : public Arbol * @param new_tfile pointer to output TFile * @param tchain pointer to TChain of input TFiles * @param branch_names std::vector of branch names to keep + * @param remove_branches bool selections so branches are dropped not kept * @return none */ - Arbusto(TFile* new_tfile, TChain* tchain, std::vector branch_names); + Arbusto(TFile* new_tfile, TChain* tchain, std::vector branch_names, bool remove_branches); /** * Arbusto object overload constructor * @param cli HEPCLI object diff --git a/rapido/src/arbusto.icc b/rapido/src/arbusto.icc index d5b9803..901b6bf 100644 --- a/rapido/src/arbusto.icc +++ b/rapido/src/arbusto.icc @@ -1,27 +1,28 @@ Arbusto::Arbusto() {} -Arbusto::Arbusto(TFile* new_tfile, TChain* tchain, std::vector keep_branch_names) -: keep_branch_names(keep_branch_names) +Arbusto::Arbusto(TFile* new_tfile, TChain* tchain, std::vector branch_names, bool remove_branches) +: remove_branches(remove_branches), branch_names(branch_names) { + // If remove_branches is true, this inverts the selection to drop the listed branches tfile = new_tfile; // Disable all branches - tchain->SetBranchStatus("*", 0); + tchain->SetBranchStatus("*", remove_branches); // Enable selected branches - for (auto branch_name : keep_branch_names) + for (auto branch_name : branch_names) { - tchain->SetBranchStatus(branch_name, 1); + tchain->SetBranchStatus(branch_name, !remove_branches); } ttree = (TTree*)tchain->CloneTree(0); } -Arbusto::Arbusto(HEPCLI& cli, std::vector keep_branch_names) -: keep_branch_names(keep_branch_names) +Arbusto::Arbusto(HEPCLI& cli, std::vector branch_names) +: branch_names(branch_names) { tfile = new TFile(TString(cli.output_dir+"/"+cli.output_name+".root"), "RECREATE"); // Disable all branches cli.input_tchain->SetBranchStatus("*", 0); // Enable selected branches - for (auto branch_name : keep_branch_names) + for (auto branch_name : branch_names) { cli.input_tchain->SetBranchStatus(branch_name, 1); } @@ -42,11 +43,11 @@ Arbusto::~Arbusto() void Arbusto::init(TTree* next_ttree) { // Disable all branches - next_ttree->SetBranchStatus("*", 0); + next_ttree->SetBranchStatus("*", remove_branches); // Enable selected branches - for (auto branch_name : keep_branch_names) + for (auto branch_name : branch_names) { - next_ttree->SetBranchStatus(branch_name, 1); + next_ttree->SetBranchStatus(branch_name, !remove_branches); } next_ttree->CopyAddresses(ttree); orig_ttree = next_ttree; diff --git a/skimmer/main.cc b/skimmer/main.cc index 486a448..3260a04 100644 --- a/skimmer/main.cc +++ b/skimmer/main.cc @@ -16,6 +16,9 @@ int main(int argc, char** argv) "RECREATE" ); + // Set to true to specify branches to DROP instead of keep + bool remove_branches = false; + // Output setting (setting which TBranches to save from original Nano) Arbusto arbusto = Arbusto( output_tfile, @@ -40,7 +43,8 @@ int main(int argc, char** argv) "SubJet*", "HLT_*", "Pileup*" - } + }, + remove_branches ); // Initialize TLists for metadata TTrees