delphes_EIC/SimpleAnalysis/SimpleAnalysis.cc

0001 #include <TROOT.h>
0002 #include <TFile.h>
0003 #include <TChain.h>
0004 #include <TTree.h>
0005 #include <TString.h>
0006 #include <TObjString.h>
0007 #include "TInterpreter.h"
0008
0009 #include <unistd.h>
0010 #include <stdlib.h>
0011 #include <iostream>
0012 #include <stdio.h>
0013 #include <getopt.h>
0014 #include <glob.h>
0015 #include <vector>
0016 #include <map>
0017 #include <any>
0018
0019 #include "classes/DelphesClasses.h"
0020 #include "external/ExRootAnalysis/ExRootTreeReader.h"
0021
0022 #include "ModuleHandler.h"
0023 #include "TreeHandler.h"
0024
0025 static std::string input_dir = "";
0026 static std::string output_file = "";
0027 static std::string module_sequence = "";
0028 static int nevents = -1;
0029
0030 ModuleHandler *ModuleHandler::instance = 0;
0031 TreeHandler *TreeHandler::instance = 0;
0032
0033
0034 // HELPER METHODS
0035
0036 void PrintHelp()
0037 {
0038   std::cout <<
0039     "--input_dir <i>:       Directory containing all the ROOT files you want to process\n"
0040     "--output_file <o>:     Output ROOT file to store results\n"
0041     "--module_sequence <s>: A string comma-separated list of modules to load; order is preserved in execution.\n"
0042     "--nevents <n>:         The total number of events to process, starting from the zeroth event in the input.\n"
0043     "--help:                Show this helpful message!\n";
0044   exit(1);
0045 }
0046
0047 std::vector<std::string> fileVector(const std::string& pattern){
0048   glob_t glob_result;
0049   glob(pattern.c_str(),GLOB_TILDE,NULL,&glob_result);
0050   std::vector<std::string> files;
0051   for(unsigned int i=0;i<glob_result.gl_pathc;++i){
0052     files.push_back(std::string(glob_result.gl_pathv[i]));
0053   }
0054   globfree(&glob_result);
0055   return files;
0056 }
0057
0058 // MAIN FUNCTION
0059
0060
0061 int main(int argc, char *argv[])
0062 {
0063   std::cout <<
0064     "===================== SIMPLEANALYSIS =====================" << std::endl;
0065
0066
0067   // Handle complex TTree data storage types by defining them for ROOT
0068   gInterpreter->GenerateDictionary("std::vector<std::vector<float>>","vector");
0069
0070
0071
0072   if (argc <= 1) {
0073     PrintHelp();
0074   }
0075
0076   const char* const short_opts = "i:o:h";
0077   const option long_opts[] = {
0078     {"input_dir", required_argument, nullptr, 'i'},
0079     {"output_file", required_argument, nullptr, 'o'},
0080     {"module_sequence", required_argument, nullptr, 's'},
0081     {"nevents", optional_argument, nullptr, 'n'},
0082     {"help", no_argument, nullptr, 'h'},
0083     {nullptr, no_argument, nullptr, 0}
0084   };
0085
0086   while (true)
0087     {
0088       const auto opt = getopt_long(argc, argv, short_opts, long_opts, nullptr);
0089
0090       if (-1 == opt)
0091     break;
0092
0093       switch (opt)
0094         {
0095         case 'i':
0096       input_dir = optarg;
0097       std::cout << "Input Directory: " << input_dir << std::endl;
0098       break;
0099
0100         case 'o':
0101       output_file = optarg;
0102       std::cout << "Output File: " << output_file << std::endl;
0103       break;
0104
0105         case 's':
0106       module_sequence = optarg;
0107       std::cout << "Module sequence: " << module_sequence << std::endl;
0108       break;
0109
0110         case 'n':
0111       nevents = std::stoi(optarg);
0112       std::cout << "Number of events to process: " << nevents << std::endl;
0113       break;
0114
0115
0116         case 'h': // -h or --help
0117         case '?': // Unrecognized option
0118       PrintHelp();
0119       break;
0120         default:
0121       PrintHelp();
0122       break;
0123         }
0124     }
0125
0126
0127   auto data = new TChain("Delphes");
0128
0129   auto files = fileVector(input_dir);
0130
0131   for (auto file : files)
0132     {
0133       data->Add(file.c_str());
0134     }
0135
0136   ExRootTreeReader *treeReader = new ExRootTreeReader(data);
0137   int n_entries = data->GetEntries();
0138
0139   std::cout
0140     << "The provided data set contains the following number of events: " << std::endl
0141     << n_entries
0142     << std::endl;
0143
0144   // Load object pointers
0145   TClonesArray *branchJet = treeReader->UseBranch("Jet");
0146   TClonesArray *branchElectron = treeReader->UseBranch("Electron");
0147   TClonesArray *branchPhoton = treeReader->UseBranch("EFlowPhoton");
0148   TClonesArray *branchNeutralHadron = treeReader->UseBranch("EFlowNeutralHadron");
0149   TClonesArray *branchGenJet = treeReader->UseBranch("GenJet");
0150   TClonesArray *branchGenParticle = treeReader->UseBranch("Particle");
0151   TClonesArray *branchRawTrack = treeReader->UseBranch("Track");
0152   TClonesArray *branchEFlowTrack = treeReader->UseBranch("EFlowTrack");
0153   TClonesArray *branchMET = treeReader->UseBranch("MissingET");
0154
0155
0156   // Setup the module handler
0157   ModuleHandler *module_handler = module_handler->getInstance(treeReader);
0158   auto module_list = TString(module_sequence).Tokenize(",");
0159
0160   // Setup the output storage
0161   TreeHandler *tree_handler = tree_handler->getInstance(output_file.c_str(), "tree");
0162   tree_handler->initialize();
0163
0164   for (int i = 0; i < module_list->GetEntries(); i++) {
0165     auto name = static_cast<TObjString*>(module_list->At(i))->GetString().Data();
0166     std::cout << "   Appending module " << name << std::endl;
0167     module_handler->addModule(name);
0168   }
0169
0170   for (auto module : module_handler->getModules()) {
0171     module->initialize();
0172   }
0173
0174   if (nevents < 0) {
0175     std::cout
0176       << "Processing all events in the sample..." << std::endl;
0177   } else {
0178     std::cout
0179       << "Processing "<< nevents << " events in the sample..." << std::endl;
0180   }
0181
0182   for(int i=0; i < n_entries; ++i) {
0183     // event number printout
0184     if(i%1000==0) {
0185       std::cout << "Processing Event " << i << std::endl;
0186     }
0187
0188     if (nevents >= 0 && i >= nevents)
0189       break;
0190
0191     // read the data for i-th event
0192     // data->GetEntry(i);
0193     // Load selected branches with data from specified event
0194     treeReader->ReadEntry(i);
0195
0196     std::map<std::string, std::any> DataStore;
0197
0198
0199     for (auto module : module_handler->getModules()) {
0200       module->setJets(branchJet);
0201       module->setGenJets(branchGenJet);
0202       module->setEFlowTracks(branchEFlowTrack);
0203       module->setTracks(branchRawTrack);
0204       module->setGenParticles(branchGenParticle);
0205       module->setPhotons(branchPhoton);
0206       module->setElectrons(branchElectron);
0207       module->setNeutralHadrons(branchNeutralHadron);
0208       module->setMET(branchMET);
0209
0210       bool result = module->execute(&DataStore);
0211       if (result == false)
0212     break;
0213     }
0214
0215     tree_handler->execute();
0216
0217
0218     // if (DataStore.find("CharmJets") != DataStore.end()) {
0219     //   std::vector<Jet*> charm_jets = std::any_cast<std::vector<Jet*>>(DataStore["CharmJets"]);
0220     // }
0221
0222   }
0223
0224   for (auto module : module_handler->getModules()) {
0225     module->finalize();
0226   }
0227   tree_handler->finalize();
0228
0229
0230   std::cout <<
0231     "========================== FINIS =========================" << std::endl;
0232
0233   exit(EXIT_SUCCESS);
0234 }