snippets/Tree_Pruning/TreePrune.C

0001 // Stephen JD Kay, University of York, 21/02/25
0002 // A short script to read in a file and prune it to only retain a smaller subset of branches.
0003 // This could be utilised to trim down a full EICrecon file for a new user to look at. This avoids the potentially overwhelming number of branches normally stored
0004 #include <string>
0005
0006 void TreePrune(TString infile=""){
0007
0008   // If no input file provide as argument, promot for one
0009   if(infile == ""){
0010     cout << "Enter a filename to analyse: ";
0011     cin >> infile;
0012   }
0013
0014   // Check input file exists, exit if not
0015   if(gSystem->AccessPathName(infile) == kTRUE){
0016     cerr << "!!!!! ERROR !!!!!" << endl << infile << " not found" << endl << "!!!!! ERROR !!!!!" << endl;
0017     exit(0);
0018   }
0019
0020   TString ofile_name, ofile_tmp, BeamE;
0021   TObjArray *tmp_Name_arr;
0022
0023   // Check the input file is a .root file as we would expect
0024   if(infile.Contains(".root") == false){
0025     cerr << "!!!!!!!!!!! - ERROR - !!!!!!!!!!!!" << endl;
0026     cerr << "Input files should be a root file!" << endl;
0027     cerr << "!!!!!!!!!!! - ERROR - !!!!!!!!!!!!" << endl;
0028     exit(1);
0029   }
0030   else{
0031     cout << "Opening and pruning " << infile << endl;
0032   }
0033
0034   // If the input file name is a file path containing /, extract only the actual file name for further use. Assign the temporary name of the output to be the input, minus its .root extension
0035   if(infile.Contains("/")){
0036     tmp_Name_arr = infile.Tokenize("/");
0037     ofile_tmp = (((TObjString *)(tmp_Name_arr->At(tmp_Name_arr->GetLast())))->String()).ReplaceAll(".root","");
0038   }
0039   else{
0040     ofile_tmp = infile;
0041     ofile_tmp.ReplaceAll(".root","");
0042   }
0043   // Set output file name
0044   ofile_name = Form("%s_Pruned.root", ofile_tmp.Data());
0045
0046   // Open our full, unpruned file
0047   TFile *full_file = TFile::Open(infile);
0048   TTree* full_tree;
0049   // Get the events tree
0050   full_file->GetObject("events", full_tree);
0051   // Deactivate all branches
0052   full_tree->SetBranchStatus("*", 0);
0053   // Activate only the branches we want to keep, add more as needed via the line below
0054   // full_tree->SetBranchStatus("",1)
0055   full_tree->SetBranchStatus("MCParticles*",1);
0056   full_tree->SetBranchStatus("ReconstructedChargedParticles*",1);
0057   full_tree->SetBranchStatus("ReconstructedChargedParticleAssociations*",1);
0058   // Note that the wild card after the branch name is needed to ensure all leaves are retained adequately. You could only retain specific leaves if you wanted to though e.g. MCParticles.PDG
0059
0060   // Set and open output file for the histograms
0061   TFile *ofile = TFile::Open(ofile_name,"RECREATE");
0062   // Clone the tree
0063   TTree* pruned_tree = full_tree->CloneTree(0);
0064   // Copy the branches
0065   pruned_tree->CopyEntries(full_tree);
0066
0067   // Write the new tree to file
0068   ofile->Write();
0069   ofile->Close(); // Close output file
0070   full_file->Close(); // Close input file
0071
0072   delete ofile;
0073   delete full_file;
0074
0075 }