rich/src/benchmark.cc

0001 // Copyright 2023, Christopher Dilks
0002 // Subject to the terms in the LICENSE file found in the top-level directory.
0003
0004 #include <functional>
0005 #include <iostream>
0006 #include <unistd.h>
0007 #include <spdlog/spdlog.h>
0008
0009 #include <TFile.h>
0010
0011 #include <podio/podioVersion.h>
0012 #if podio_VERSION >= PODIO_VERSION(0, 99, 0)
0013 #include <podio/ROOTReader.h>
0014 #else
0015 #include <podio/ROOTFrameReader.h>
0016 #endif
0017 #include <podio/Frame.h>
0018
0019 #include "SimHitAnalysis.h"
0020 #include "RawHitAnalysis.h"
0021 #include "CherenkovPIDAnalysis.h"
0022 #include "ReconstructedParticleAnalysis.h"
0023
0024 using namespace benchmarks;
0025
0026 // -------------------------------------------------------------
0027
0028 // main logger
0029 std::shared_ptr<spdlog::logger> m_log;
0030
0031 // get a collection, and check its validity
0032 template<class C>
0033 const C& GetCollection(podio::Frame& frame, std::string name) {
0034   const auto& coll = frame.get<C>(name);
0035   if(!coll.isValid())
0036     m_log->error("invalid collection '{}'", name);
0037   return coll;
0038 }
0039
0040 // -------------------------------------------------------------
0041 // MAIN
0042 // -------------------------------------------------------------
0043 int main(int argc, char** argv) {
0044
0045   // -------------------------------------------------------------
0046   // setup and CLI option parsing
0047   // -------------------------------------------------------------
0048
0049   // loggers
0050   m_log = spdlog::default_logger()->clone("benchmark_rich");
0051   auto log_level = spdlog::level::info;
0052   m_log->set_level(log_level);
0053
0054   // default options
0055   int verbosity             = 0;
0056   long long num_events      = 0;
0057   std::string ana_file_name = "out_rich.root";
0058   std::vector<std::string> rec_files;
0059   std::vector<std::string> analysis_algorithms = {
0060     "SimHit",
0061     "RawHit",
0062     "CherenkovPID",
0063     "ReconstructedParticle"
0064   };
0065   auto list_algorithms = [&analysis_algorithms] () {
0066     for(auto analysis_algorithm : analysis_algorithms)
0067       std::cout << "        " << analysis_algorithm << std::endl;
0068     std::cout << std::endl;
0069   };
0070
0071   // usage guide
0072   auto PrintUsage = [&argv, &list_algorithms, ana_file_name] () {
0073     std::cout << "\nUSAGE: " << argv[0] << " -i [INPUT FILES]... [OPTIONS]..." << std::endl
0074       << "\nREQUIRED OPTIONS:\n"
0075       << " -i [INPUT FILES]...\n"
0076       << "    input file from reconstruction output;\n"
0077       << "    can specify more than one (delimited by space)\n"
0078       << "\nOPTIONAL OPTIONS:\n"
0079       << " -o [OUTPUT FILE]\n"
0080       << "    output analysis file (default: " << ana_file_name << ")\n"
0081       << " -n [NUMBER OF EVENTS]\n"
0082       << "    number of events to process (default: process all)\n"
0083       << " -v: verbose output\n"
0084       << " -vv: more verbose output\n"
0085       << " -a [ALGORITHMS]...\n"
0086       << "    list of analysis algorithms to run (delimited by spaces)\n"
0087       << "    default: run all of them\n"
0088       << "    available algorithms:\n";
0089     list_algorithms();
0090   };
0091   if(argc==1) {
0092     PrintUsage();
0093     return 2;
0094   }
0095
0096   // parse options
0097   int opt;
0098   while( (opt=getopt(argc, argv, "i:o:n:v|a:")) != -1) {
0099     switch(opt) {
0100       case 'i':
0101         optind--;
0102         for(; optind<argc && *argv[optind]!='-'; optind++)
0103           rec_files.push_back(std::string(argv[optind]));
0104         break;
0105       case 'o':
0106         ana_file_name = std::string(optarg);
0107         break;
0108       case 'n':
0109         num_events = std::atoll(optarg);
0110         break;
0111       case 'v':
0112         verbosity++;
0113         break;
0114       case 'a':
0115         analysis_algorithms.clear();
0116         optind--;
0117         for(; optind<argc && *argv[optind]!='-'; optind++)
0118           analysis_algorithms.push_back(std::string(argv[optind]));
0119         break;
0120       default:
0121         m_log->error("Unknown option '{}'", opt);
0122         PrintUsage();
0123         return 1;
0124     }
0125   }
0126
0127   // print options
0128   m_log->info("{:-^50}", " User options ");
0129   m_log->info("  input files:");
0130   for(auto rec_file : rec_files)
0131     m_log->info("    {}", rec_file);
0132   m_log->info("  output file: {}", ana_file_name);
0133   m_log->info("  algorithms:");
0134   for(auto analysis_algorithm : analysis_algorithms)
0135     m_log->info("    {}", analysis_algorithm);
0136   m_log->info("  num_events: {}", num_events);
0137   m_log->info("  verbosity: {}", verbosity);
0138   m_log->info("{:-^50}", "");
0139
0140   // check options
0141   std::vector<std::string> bad_options;
0142   if(rec_files.size()==0) bad_options.push_back("no [INPUT FILES] have been specified");
0143   if(analysis_algorithms.size()==0) bad_options.push_back("no [ALGORITHMS] have been specified");
0144   for(auto bad_option : bad_options) m_log->error(bad_option);
0145   if(bad_options.size()>0) {
0146     PrintUsage();
0147     return 1;
0148   }
0149
0150   // set log level
0151   switch(verbosity) {
0152     case 0:  log_level = spdlog::level::info;  break;
0153     case 1:  log_level = spdlog::level::debug; break;
0154     default: log_level = spdlog::level::trace;
0155   }
0156   m_log->set_level(log_level);
0157
0158   // start the output file
0159   auto ana_file = new TFile(TString(ana_file_name), "RECREATE");
0160
0161   // struct to hold an algorithm execution processor
0162   struct AnalysisAlgorithm {
0163     std::shared_ptr<spdlog::logger>    log;     // algorithm-specific logger
0164     std::function<void(podio::Frame&)> process; // call AlgorithmProcess
0165     std::function<void()>              finish;  // call AlgorithmFinish
0166   };
0167   std::vector<AnalysisAlgorithm> algorithm_processors;
0168
0169   // loop over algorithms to run, and define how to run them
0170   for(auto analysis_algorithm : analysis_algorithms) {
0171     AnalysisAlgorithm algo;
0172
0173     // algorithm logger
0174     algo.log = spdlog::default_logger()->clone(analysis_algorithm);
0175     algo.log->set_level(log_level);
0176     algo.log->debug("{:-<50}","INITIALIZE ");
0177
0178     // --------------------------------------------------------------
0179     // define how to run each algorithm
0180     // --------------------------------------------------------------
0181
0182     // truth hits (photons) .........................................
0183     if(analysis_algorithm == "SimHit") {
0184       auto sim_algo = std::make_shared<SimHitAnalysis>();
0185       ana_file->mkdir("phot")->cd();
0186       sim_algo->AlgorithmInit(algo.log);
0187       algo.process = [sim_algo] (podio::Frame& frame) {
0188         sim_algo->AlgorithmProcess(
0189             GetCollection<edm4hep::MCParticleCollection>(frame,"MCParticles"),
0190             GetCollection<edm4hep::SimTrackerHitCollection>(frame,"DRICHHits")
0191             );
0192       };
0193       algo.finish = [sim_algo] () { sim_algo->AlgorithmFinish(); };
0194     }
0195
0196     // digitizer ....................................................
0197     else if(analysis_algorithm == "RawHit") {
0198       auto digi_algo = std::make_shared<RawHitAnalysis>();
0199       ana_file->mkdir("digi")->cd();
0200       digi_algo->AlgorithmInit(algo.log);
0201       algo.process = [digi_algo] (podio::Frame& frame) {
0202         digi_algo->AlgorithmProcess(
0203             GetCollection<edm4eic::RawTrackerHitCollection>(frame,"DRICHRawHits"),
0204             GetCollection<edm4eic::MCRecoTrackerHitAssociationCollection>(frame,"DRICHRawHitsAssociations")
0205             );
0206       };
0207       algo.finish = [digi_algo] () { digi_algo->AlgorithmFinish(); };
0208     }
0209
0210     // IRT ..........................................................
0211     else if(analysis_algorithm == "CherenkovPID") {
0212       std::vector<std::string> radiators = { "Aerogel", "Gas", "Merged" };
0213       std::map<std::string, std::shared_ptr<CherenkovPIDAnalysis>> irt_algos;
0214       for(auto radiator : radiators)
0215         irt_algos.insert({radiator, std::make_shared<CherenkovPIDAnalysis>()});
0216
0217       for(auto& [radiator, irt_algo] : irt_algos) {
0218         ana_file->mkdir(("pid"+radiator).c_str())->cd();
0219         irt_algo->AlgorithmInit(radiator, algo.log);
0220       }
0221
0222       algo.process = [irt_algos] (podio::Frame& frame) {
0223         const auto& mc_parts = GetCollection<edm4hep::MCParticleCollection>(frame,"MCParticles");
0224         for(auto& [radiator, irt_algo] : irt_algos) {
0225           const auto& cherenkov_pids = GetCollection<edm4eic::CherenkovParticleIDCollection>(
0226               frame,
0227               "DRICH" + radiator + "IrtCherenkovParticleID"
0228               );
0229           irt_algo->AlgorithmProcess(mc_parts, cherenkov_pids);
0230         }
0231       };
0232
0233       algo.finish = [irt_algos] () {
0234         for(auto& [radiator, irt_algo] : irt_algos)
0235           irt_algo->AlgorithmFinish();
0236       };
0237
0238     }
0239
0240     // linking to reconstructed particles ...........................
0241     else if(analysis_algorithm == "ReconstructedParticle") {
0242       auto link_algo = std::make_shared<ReconstructedParticleAnalysis>();
0243       ana_file->mkdir("link")->cd();
0244       link_algo->AlgorithmInit(algo.log);
0245       algo.process = [link_algo] (podio::Frame& frame) {
0246         link_algo->AlgorithmProcess(
0247             GetCollection<edm4eic::MCRecoParticleAssociationCollection>(
0248               frame,
0249               "ReconstructedChargedParticleAssociations"
0250               )
0251             );
0252       };
0253       algo.finish = [link_algo] () { link_algo->AlgorithmFinish(); };
0254     }
0255
0256     // unknown algorithm ............................................
0257     else {
0258       m_log->error("Unknown [ALGORITHM] '{}'; ignoring", analysis_algorithm);
0259       continue;
0260     }
0261
0262     // --------------------------------------------------------------
0263
0264     algorithm_processors.push_back(algo);
0265   }
0266
0267   // -------------------------------------------------------------
0268   // read the input file and run all algorithms
0269   // -------------------------------------------------------------
0270
0271   // open the input files
0272 #if podio_VERSION >= PODIO_VERSION(0, 99, 0)
0273   podio::ROOTReader podioReader;
0274 #else
0275   podio::ROOTFrameReader podioReader;
0276 #endif
0277   m_log->warn("podio::ROOTFrameReader cannot yet support multiple files; reading only the first");
0278   // podioReader.openFiles(rec_files);
0279   podioReader.openFile(rec_files.front());
0280
0281   // get the number of events to process
0282   const std::string tree_name = "events";
0283   long long num_entries = podioReader.getEntries(tree_name);
0284   if(num_events>0) num_entries = std::min(num_events, num_entries);
0285
0286   // event loop
0287   m_log->info("BEGIN EVENT LOOP");
0288   for(long long e=0; e<num_entries; e++) {
0289     if(e%100==0) m_log->info("Progress: {:.3f}%", 100.0 * e/num_entries);
0290     auto podioFrame = podio::Frame(podioReader.readNextEntry(tree_name));
0291     for(auto& algo : algorithm_processors) {
0292       algo.log->debug("{:-<50}","PROCESS EVENT ");
0293       algo.process(podioFrame);
0294     }
0295   }
0296
0297   // finish
0298   for(auto& algo : algorithm_processors) {
0299     algo.log->debug("{:-<50}","FINISH ");
0300     algo.finish();
0301   }
0302
0303   // write output
0304   ana_file->Write();
0305   ana_file->Close();
0306   m_log->info("Done. Wrote analysis file:");
0307   m_log->info("  {}", ana_file_name);
0308 }