tutorial-analysis/files/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("ReconstructedParticles**",1);
0057   full_tree->SetBranchStatus("ReconstructedChargedParticles*",1);
0058   full_tree->SetBranchStatus("ReconstructedParticleAssociations*",1);
0059   full_tree->SetBranchStatus("ReconstructedChargedParticleAssociations*",1);
0060   full_tree->SetBranchStatus("EcalEndcapNClusters*",1);
0061   full_tree->SetBranchStatus("EcalEndcapPClusters*",1);
0062
0063   // Set and open output file for the histograms
0064   TFile *ofile = TFile::Open(ofile_name,"RECREATE");
0065   // Clone the tree
0066   TTree* pruned_tree = full_tree->CloneTree(0);
0067   // Copy the branches
0068   pruned_tree->CopyEntries(full_tree);
0069
0070   // Write the new tree to file
0071   ofile->Write();
0072   ofile->Close(); // Close output file
0073   full_file->Close(); // Close input file
0074
0075   delete ofile;
0076   delete full_file;
0077
0078 }