Merge pull request #14 from ireed3597/master

Add ability to specify which branches to drop instead of keep
cmstas · Nov 2, 2023 · 9f36748 · 9f36748
2 parents a8d83c4 + 71e9faf
commit 9f36748
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 14 deletions.
diff --git 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<TString> keep_branch_names;
+    std::vector<TString> 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<TString> branch_names);
+    Arbusto(TFile* new_tfile, TChain* tchain, std::vector<TString> 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
@@ -1,27 +1,28 @@
 Arbusto::Arbusto() {}
 
-Arbusto::Arbusto(TFile* new_tfile, TChain* tchain, std::vector<TString> keep_branch_names) 
-: keep_branch_names(keep_branch_names)
+Arbusto::Arbusto(TFile* new_tfile, TChain* tchain, std::vector<TString> 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<TString> keep_branch_names) 
-: keep_branch_names(keep_branch_names)
+Arbusto::Arbusto(HEPCLI& cli, std::vector<TString> 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
@@ -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