Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-07-11 07:50:13

0001 #ifdef __MACH__
0002 #include <mach/mach.h>
0003 #endif
0004 
0005 #include <iostream>
0006 #include <fstream>
0007 #include <string>
0008 #include <random>
0009 #include <ctime>
0010 #include <cstdlib>
0011 #include <vector>
0012 #include <algorithm>
0013 #include <numeric>
0014 #include <stdexcept>
0015 #include <chrono>
0016 #include <cmath>
0017 #include <random>
0018 #include <tuple>
0019 #include <unistd.h>
0020 #include <sys/resource.h>
0021 
0022 #include <HepMC3/ReaderFactory.h>
0023 #include <HepMC3/WriterAscii.h>
0024 #include "HepMC3/WriterRootTree.h"
0025 #include "HepMC3/GenRunInfo.h"
0026 #include "HepMC3/GenEvent.h"
0027 #include "HepMC3/Print.h"
0028 
0029 #include "argparse/argparse.hpp"
0030 
0031 using std::cout;
0032 using std::cerr;
0033 using std::endl;
0034 using std::string;
0035 
0036 #define _LITERAL_TO_STRING(s) #s
0037 #define _AS_STRING(s) _LITERAL_TO_STRING(s)
0038 const char* hepmc_merger_version = _AS_STRING(HEPMC_MERGER_VERSION_FULL);
0039 #undef _LITERAL_TO_STRING
0040 #undef _AS_STRING
0041 
0042 // =============================================================
0043 /**
0044     Combine signal and up to four background HEPMC files.
0045     
0046     Typical usage:
0047     ./SignalBackgroundMerger --signalFile dis.hepmc3.tree.root --signalFreq 0 \
0048             --bgFile hgas.hepmc3.tree.root 2000 0 2000 \
0049         --bgFile egastouschk.hepmc3.tree.root 20 0 3000 \
0050             --bgFile egascouloumb.hepmc3.tree.root 20 0 4000 \
0051         --bgFile egasbrems.hepmc3.tree.root 20 0 5000 \
0052             --bgFile synrad.hepmc3.tree.root 25 0 6000
0053 **/    
0054 
0055 struct BackgroundConfig {
0056     std::string file;
0057     double frequency=0;
0058     int skip=0;
0059     int status=0;
0060 } ;
0061 
0062 class SignalBackgroundMerger {
0063 
0064 private:
0065   // more private data at the end; pulling these more complicated objects up for readability
0066   std::shared_ptr<HepMC3::Reader> sigAdapter;
0067   double sigFreq = 0;
0068   int sigStatus = 0;
0069   std::map<std::string, std::shared_ptr<HepMC3::Reader>> freqAdapters;
0070   std::map<std::string, double> freqs;
0071   std::map<std::string, int> baseStatuses;
0072 
0073   std::map<std::string,
0074       std::tuple<std::vector<HepMC3::GenEvent>,
0075               std::piecewise_constant_distribution<>,
0076               double>
0077        > weightDict;
0078 
0079   // just keep count of some numbers, could be more sophisticated
0080   typedef struct{
0081     long eventCount;
0082     long particleCount;
0083   } stats;
0084   std::map<std::string, stats > infoDict;
0085 
0086 public:
0087 
0088   SignalBackgroundMerger(int argc, char* argv[]) {
0089     auto t0 = std::chrono::high_resolution_clock::now();
0090 
0091     // Parse arguments, print banner, open files, initialize rng
0092     digestArgs(argc, argv);
0093     rng.seed( rngSeed );
0094     banner();
0095     if (outputFile != "" ) {
0096       outputFileName = outputFile;
0097     } else {
0098       outputFileName = nameGen();
0099     }
0100     std::cout << "\n==================================================================\n";
0101     cout << "Writing to " << outputFileName << endl;
0102 
0103     PrepData ( signalFile, signalFreq, signalSkip, signalStatus, true );
0104     for (const auto& bg : backgroundFiles) {
0105     PrepData ( bg.file, bg.frequency, bg.skip, bg.status, false );
0106     }
0107     
0108     
0109     auto t1 = std::chrono::high_resolution_clock::now();
0110     std::cout << "Initiation time: " << std::round(std::chrono::duration<double, std::chrono::seconds::period>(t1 - t0).count()) << " sec" << std::endl;
0111     std::cout << "\n==================================================================\n" << std::endl;
0112 
0113   }
0114 
0115   // Helper to parse raw strings into BackgroundConfig structs
0116   std::vector<BackgroundConfig>
0117   parse_backgrounds(const std::vector<std::string> &raw_args_list) {
0118     std::vector<BackgroundConfig> backgrounds;
0119     auto is_pure_integer = [](const std::string &str) {
0120       if (str.empty()) return false;
0121       for (char c : str) {
0122         if (!std::isdigit(c)) return false;
0123       }
0124       return true;
0125     };
0126       
0127     // Group strings into sets of 2-4 arguments per background
0128     for (size_t i = 0; i < raw_args_list.size();) {
0129       // Determine how many arguments this background has
0130       size_t args_count = 2; // minimum
0131 
0132       // Look ahead to see if next strings can be parsed as numbers (skip/status)
0133       if (i + 2 < raw_args_list.size() && is_pure_integer(raw_args_list[i + 2])) {
0134           args_count = 3;
0135 
0136           if (i + 3 < raw_args_list.size() && is_pure_integer(raw_args_list[i + 3])) {
0137               args_count = 4;
0138           }
0139       }
0140 
0141       // Ensure we don't go beyond the vector bounds
0142       if (i + args_count > raw_args_list.size()) {
0143         args_count = raw_args_list.size() - i;
0144       }
0145 
0146       if (args_count < 2) {
0147         throw std::runtime_error("Background file " +
0148                                  std::to_string(backgrounds.size()) +
0149                                  " must have at least 2 arguments");
0150       }
0151 
0152       try {
0153         BackgroundConfig bg;
0154         bg.file = raw_args_list[i];
0155         bg.frequency = std::stod(raw_args_list[i + 1]);
0156         bg.skip = (args_count > 2) ? std::stoi(raw_args_list[i + 2]) : 0;
0157         bg.status = (args_count > 3) ? std::stoi(raw_args_list[i + 3]) : 0;
0158         backgrounds.push_back(bg);
0159       } catch (const std::exception &e) {
0160         throw std::runtime_error("Error parsing background file " +
0161                                  std::to_string(backgrounds.size()) + ": " +
0162                                  e.what());
0163       }
0164 
0165       i += args_count;
0166     }
0167 
0168     return backgrounds;
0169   }
0170 
0171   void merge(){
0172     auto t1 = std::chrono::high_resolution_clock::now();
0173 
0174     // Populate run-level metadata
0175     auto runInfo = std::make_shared<HepMC3::GenRunInfo>();
0176 
0177     runInfo->add_attribute("hepmc_merger_version",
0178         std::make_shared<HepMC3::StringAttribute>(hepmc_merger_version));
0179     runInfo->add_attribute("hepmc_merger_signal_file",
0180         std::make_shared<HepMC3::StringAttribute>(signalFile));
0181     runInfo->add_attribute("hepmc_merger_signal_frequency_kHz",
0182         std::make_shared<HepMC3::DoubleAttribute>(signalFreq));
0183     runInfo->add_attribute("hepmc_merger_signal_skip",
0184         std::make_shared<HepMC3::IntAttribute>(signalSkip));
0185 
0186     std::string bgFiles, bgFreqs, bgAvgRates, bgSkips;
0187     for (size_t bi = 0; bi < backgroundFiles.size(); ++bi) {
0188       if (bi) { bgFiles += ";"; bgFreqs += ";"; bgAvgRates += ";"; bgSkips += ";"; }
0189       bgFiles += backgroundFiles[bi].file;
0190       bgFreqs += std::to_string(backgroundFiles[bi].frequency);
0191       bgSkips += std::to_string(backgroundFiles[bi].skip);
0192       double avgRate = 0.0;
0193       if (backgroundFiles[bi].frequency <= 0.0) {
0194         auto it = weightDict.find(backgroundFiles[bi].file);
0195         if (it != weightDict.end())
0196           avgRate = std::get<2>(it->second) * 1e6; // GHz -> kHz
0197       }
0198       bgAvgRates += std::to_string(avgRate);
0199     }
0200     runInfo->add_attribute("hepmc_merger_background_files",
0201         std::make_shared<HepMC3::StringAttribute>(bgFiles));
0202     runInfo->add_attribute("hepmc_merger_background_frequencies_kHz",
0203         std::make_shared<HepMC3::StringAttribute>(bgFreqs));
0204     runInfo->add_attribute("hepmc_merger_background_avg_rates_kHz",
0205         std::make_shared<HepMC3::StringAttribute>(bgAvgRates));
0206     runInfo->add_attribute("hepmc_merger_background_skips",
0207         std::make_shared<HepMC3::StringAttribute>(bgSkips));
0208 
0209     runInfo->add_attribute("hepmc_merger_integration_window_ns",
0210         std::make_shared<HepMC3::DoubleAttribute>(intWindow));
0211     runInfo->add_attribute("hepmc_merger_n_slices",
0212         std::make_shared<HepMC3::IntAttribute>(nSlices));
0213 
0214     // Open output file — pass runInfo to constructor so it is written to the header
0215     std::shared_ptr<HepMC3::Writer> f;
0216     if (rootFormat)
0217       f = std::make_shared<HepMC3::WriterRootTree>(outputFileName, runInfo);
0218     else
0219       f = std::make_shared<HepMC3::WriterAscii>(outputFileName, runInfo);
0220 
0221     // Slice loop
0222     int i = 0;
0223     for (i = 0; i<nSlices; ++i ) {
0224       if (i % 1000 == 0 || verbose ) squawk(i);
0225       auto hepSlice = mergeSlice(i);
0226       if (!hepSlice) {
0227     std::cout << "Exhausted signal source." << std::endl;
0228     break;
0229       }
0230       hepSlice->set_event_number(i);
0231       f->write_event(*hepSlice);
0232     }
0233     std::cout << "Finished all requested slices." << std::endl;
0234 
0235     int slicesDone = i;
0236     auto t2 = std::chrono::high_resolution_clock::now();
0237 
0238     std::cout << "Slice loop time: " << std::round(std::chrono::duration<double, std::chrono::minutes::period>(t2 - t1).count()) << " min" << std::endl;
0239     std::cout << " -- " << std::round(std::chrono::duration<double, std::chrono::microseconds::period>(t2 - t1).count() / i) << " us / slice" << std::endl;
0240 
0241     for (auto info : infoDict) {
0242       std::cout << "From " << info.first << std::endl;
0243       std::cout << "  placed " << info.second.eventCount << " events" << std::endl; 
0244       std::cout << "  --> on average " << std::setprecision(3) << info.second.eventCount / float(nSlices) << std::endl;
0245       std::cout << "  placed " << info.second.particleCount << " final state particles" << std::endl;
0246       std::cout << "  --> on average " << std::setprecision(3) << info.second.particleCount / float(nSlices) << std::endl;
0247       
0248     }
0249 
0250     struct rusage r_usage;
0251     getrusage(RUSAGE_SELF, &r_usage);
0252 
0253     // NOTE: Reported in kB on Linux, bytes in Mac/Darwin
0254     // Could try to explicitly catch __linux__ as well
0255     // Unclear in BSD, I've seen conflicting reports
0256 #ifdef __MACH__
0257     float mbsize = 1024 * 1024;
0258 #else // Linux
0259     float mbsize = 1024;
0260 #endif
0261   
0262 
0263     std::cout << endl << "Maximum Resident Memory " << r_usage.ru_maxrss / mbsize << " MB" << std::endl;
0264     // clean up, close all files
0265     sigAdapter->close();
0266     for (auto& it : freqAdapters) {
0267       it.second->close();
0268     }
0269     f->close();
0270     
0271   }
0272   
0273   // ---------------------------------------------------------------------------
0274   void digestArgs(int argc, char* argv[]) {
0275     // Handle the command line tedium
0276     // ArgumentParser is meant to be used in a single function.
0277     // ArgumentParser internally uses std::string_views,
0278     // references, iterators, etc.
0279     // Many of these elements become invalidated after a copy or move.
0280     argparse::ArgumentParser args ("Merge signal events with up to four background sources.", hepmc_merger_version);
0281     
0282     args.add_argument("-i", "--signalFile")
0283       .default_value(std::string("root://dtn-eic.jlab.org//volatile/eic/EPIC/EVGEN/SIDIS/pythia6-eic/1.0.0/10x100/q2_0to1/pythia_ep_noradcor_10x100_q2_0.000000001_1.0_run1.ab.hepmc3.tree.root"))
0284       .help("Name of the HEPMC file with the signal events");
0285     
0286     args.add_argument("-sf", "--signalFreq")
0287       .default_value(0.0)
0288       .scan<'g', double>()
0289       .help("Signal frequency in kHz. Default is 0 to have exactly one signal event per slice. Set to the estimated DIS rate to randomize.");
0290 
0291     args.add_argument("-S", "--signalSkip")
0292       .default_value(0)
0293     .scan<'i', int>()
0294     .help("Number of signals events to skip. Default is 0.");
0295 
0296     args.add_argument("-St", "--signalStatus")
0297       .default_value(0)
0298     .scan<'i', int>()
0299     .help("Apply shift on particle generatorStatus code for signal. Default is 0. ");
0300 
0301     args.add_argument("-b","--bgFile")
0302       .nargs(2,4)
0303       .append()
0304       .help("Tuple with name of the HEPMC file with background events, background frequency in kHz, number of background events to skip (default 0), shift on particle generatorStatus code (default 0).");
0305     
0306     args.add_argument("-o", "--outputFile")
0307       .default_value(std::string("bgmerged.hepmc3.tree.root"))
0308       .help("Specify the output file name. By default bgmerged.hepmc3.tree.root is used");
0309 
0310     args.add_argument("-r", "--rootFormat")
0311       .default_value(true)
0312       .implicit_value(true)
0313       .help("Use hepmc.root output format, default is true.");
0314     
0315     args.add_argument("-w", "--intWindow")
0316       .default_value(2000.0)
0317       .scan<'g', double>()
0318       .help("Length of the integration window in nanoseconds. Default is 2000.");
0319     
0320     args.add_argument("-N", "--nSlices")
0321       .default_value(10000)
0322       .scan<'i', int>()
0323       .help("Number of sampled time slices ('events'). Default is 10000. If set to -1, all events in the signal file will be used and background files cycled as needed.");
0324     
0325     args.add_argument("--squashTime")
0326       .default_value(false)
0327       .implicit_value(true)
0328       .help("Integration is performed but no time information is associated to vertices.");
0329     
0330     args.add_argument("--rngSeed")
0331       .default_value(0)
0332       .action([](const std::string& value) { return std::stoi(value); })
0333       .help("Random seed, default is None");
0334     
0335     args.add_argument("-v", "--verbose")
0336       .default_value(false)
0337       .implicit_value(true)
0338       .help("Display details for every slice.");
0339 
0340     try {
0341       args.parse_args(argc, argv);
0342     }
0343     catch (const std::runtime_error& err) {
0344       std::cout << err.what() << std::endl;
0345       std::cout << args;
0346       exit(EXIT_FAILURE);
0347     }
0348     // Access arguments using args.get method
0349     signalFile = args.get<std::string>("--signalFile");
0350     signalFreq = args.get<double>("--signalFreq");
0351     signalSkip = args.get<int>("--signalSkip");
0352     signalStatus = args.get<int>("--signalStatus");
0353     backgroundFiles = parse_backgrounds(args.get<std::vector<std::string>>("--bgFile"));
0354     outputFile = args.get<std::string>("--outputFile");
0355     rootFormat = args.get<bool>("--rootFormat");
0356     intWindow  = args.get<double>("--intWindow");
0357     nSlices    = args.get<int>("--nSlices");
0358     squashTime = args.get<bool>("--squashTime");
0359     rngSeed    = args.get<int>("--rngSeed");
0360     verbose    = args.get<bool>("--verbose");
0361 
0362     
0363   }
0364   
0365   // ---------------------------------------------------------------------------
0366   void banner() {
0367     // Print banner
0368     std::cout << "==================================================================" << std::endl;
0369     std::cout << "=== EPIC HEPMC MERGER ===" << std::endl;
0370     std::cout << "authors: Benjamen Sterwerf* (bsterwerf@berkeley.edu), Kolja Kauder** (kkauder@bnl.gov), Reynier Cruz-Torres***" << std::endl;
0371     std::cout << "* University of California, Berkeley" << std::endl;
0372     std::cout << "** Brookhaven National Laboratory" << std::endl;
0373     std::cout << "*** formerly Lawrence Berkeley National Laboratory" << std::endl;
0374     std::cout << "\nFor more information, run \n./signal_background_merger --help" << std::endl;
0375 
0376     std::string statusMessage = "Shifting all particle status codes from this source by ";
0377     std::vector<int> statusList_stable, statusList_decay;
0378 
0379     std::cout << "Number of Slices:" << nSlices << endl;
0380     std::string freqTerm = signalFreq > 0 ? std::to_string(signalFreq) + " kHz" : "(one event per time slice)";
0381     std::string statusTerm = signalStatus > 0 ? statusMessage + std::to_string(signalStatus): "";
0382     if (signalStatus>0){
0383       statusList_stable.push_back(signalStatus+1);
0384       statusList_decay.push_back(signalStatus+2);
0385     }
0386     std::cout << "Signal events file and frequency:\n";
0387     std::cout << "\t- " << signalFile << "\t" << freqTerm << "\n" << statusTerm << "\n";
0388     
0389     std::cout << "\nBackground files and their respective frequencies:\n";
0390 
0391     for (const auto& bg : backgroundFiles) {
0392       if (!bg.file.empty()) {
0393         freqTerm = bg.frequency > 0 ? std::to_string(bg.frequency) + " kHz" : "(from weights)";
0394         statusTerm = bg.status > 0 ? statusMessage + std::to_string(bg.status) : "";
0395         std::cout << "\t- " << bg.file << "\t" << freqTerm << "\n" << statusTerm << "\n";
0396         if (bg.status>0){
0397           statusList_stable.push_back(bg.status+1);
0398           statusList_decay.push_back(bg.status+2);
0399         }
0400       }
0401     }
0402     
0403     auto join = [](const std::vector<int>& vec) {
0404         return std::accumulate(vec.begin(), vec.end(), std::string(),
0405             [](const std::string& a, int b) {
0406                 return a.empty() ? std::to_string(b) : a + " " + std::to_string(b);
0407             });
0408     };
0409     std::string stableStatuses = join(statusList_stable);
0410     std::string decayStatuses = join(statusList_decay);
0411     std::string message = "\n!!!Attention!!!\n To proceed the shifted particles statuses in DD4hep, please add the following options to ddsim:\n"
0412                           "--physics.alternativeStableStatuses=\"" + stableStatuses + 
0413                           "\"  --physics.alternativeDecayStatuses=\"" + decayStatuses + "\"\n";
0414     std::cout << message<<std::endl;           
0415   }
0416   
0417   // ---------------------------------------------------------------------------  
0418   void PrepData(const std::string& fileName, double freq, int skip=0, int baseStatus=0, bool signal=false) {
0419     if (fileName.empty()) return;
0420 
0421     cout << "Prepping " << fileName << endl;
0422     std::shared_ptr<HepMC3::Reader> adapter;
0423     try {
0424       adapter = HepMC3::deduce_reader(fileName);
0425       if (!adapter) {
0426         throw std::runtime_error("Failed to open file");
0427       }
0428     } catch (const std::runtime_error& e) {
0429       std::cerr << "Opening " << fileName << " failed: " << e.what() << std::endl;
0430       exit(EXIT_FAILURE);
0431     }
0432     
0433     infoDict[fileName] = {0,0};
0434 
0435     if (signal) {
0436       sigAdapter = adapter;
0437       sigFreq = freq;
0438       sigStatus = baseStatus;
0439       sigAdapter->skip(skip);
0440       return;
0441     }
0442 
0443     // Now catch the weighted case
0444     if (freq <= 0) {
0445       std::cout << "Reading in all events from " << fileName << std::endl;
0446       std::vector<HepMC3::GenEvent> events;
0447       std::vector<double> weights;
0448 
0449       while(!adapter->failed()) {
0450         HepMC3::GenEvent evt(HepMC3::Units::GEV,HepMC3::Units::MM);
0451         adapter->read_event(evt);
0452 
0453         // remove events with 0 weight - note that this does change avgRate = <weight> (by a little)
0454         if (double w=evt.weight() > 0){
0455           events.push_back(evt);
0456           weights.push_back(evt.weight());
0457           }
0458       }
0459       adapter->close();
0460       
0461       double avgRate = 0.0;
0462       for ( auto w : weights ){ avgRate += w;}
0463       avgRate /= weights.size();
0464       avgRate *= 1e-9; // convert to 1/ns == GHz
0465       std::cout << "Average rate is " << avgRate << " GHz" << std::endl;
0466 
0467       std::vector<int> indices (weights.size());
0468       std::iota (std::begin(indices), std::end(indices), 0); // [ 0 , ... , N ] <- draw randomly from this
0469       
0470       // Replacing python's compact toPlace = self.rng.choice( a=events, size=nEvents, p=probs, replace=False )
0471       // is tricky. Possibly more elegant or faster versions exist,
0472       // https://stackoverflow.com/questions/42926209/equivalent-function-to-numpy-random-choice-in-c
0473       // we'll do it rather bluntly, since the need for this code should go away soon with new SR infrastructure
0474       // https://stackoverflow.com/questions/1761626/weighted-random-numbers
0475       // Normalizing is not necessary for this method 
0476       // for ( auto& w : weights ) {
0477       //    w /= avgRate;
0478       // }
0479       std::piecewise_constant_distribution<> weightedDist(std::begin(indices),std::end(indices),
0480                               std::begin(weights));
0481       weightDict[fileName] = { std::make_tuple(events, weightedDist, avgRate) };
0482 
0483       return;
0484     }
0485 
0486     // Not signal and not weighted --> prepare frequency backgrounds
0487     adapter->skip(skip);
0488     freqAdapters[fileName] = adapter;
0489     freqs[fileName] = freq;
0490     baseStatuses[fileName] = baseStatus;
0491   }
0492 
0493    // ---------------------------------------------------------------------------
0494   bool hasEnding (std::string const &fullString, std::string const &ending) {
0495     if (fullString.length() >= ending.length()) {
0496       return (0 == fullString.compare (fullString.length() - ending.length(), ending.length(), ending));
0497     } else {
0498       return false;
0499     }
0500   }
0501   
0502   // ---------------------------------------------------------------------------
0503   std::string nameGen() {
0504     // Generate a name for the output file
0505     // It's too simplistic for input with directories
0506     std::string name = signalFile;
0507     if (nSlices > 0) {
0508         size_t pos = name.find(".hepmc");
0509         if (pos != std::string::npos) {
0510       name.replace(pos, 6, "_n_" + std::to_string(nSlices) + ".hepmc");
0511         }
0512     }
0513 
0514     if ( rootFormat && !hasEnding(name,".root")){
0515       name.append(".root");
0516     }
0517     name = "bgmerged_" + name;
0518 
0519     return name;
0520   }
0521 
0522   // ---------------------------------------------------------------------------
0523   void squawk(int i) {
0524 
0525     // More fine-grained info about current usage
0526 #ifdef __MACH__
0527     task_basic_info_data_t info;
0528     mach_msg_type_number_t size = sizeof(info);
0529     kern_return_t kerr = task_info(mach_task_self(),
0530                                    TASK_BASIC_INFO,
0531                                    (task_info_t)&info,
0532                                    &size);
0533 
0534     long memory_usage = -1;
0535     if (kerr == KERN_SUCCESS) {
0536       memory_usage = info.resident_size  / 1024 / 1024;
0537     }
0538 #else // Linux
0539     std::ifstream statm("/proc/self/statm");
0540     long size, resident, share, text, lib, data, dt;
0541     statm >> size >> resident >> share >> text >> lib >> data >> dt;
0542     statm.close();
0543 
0544     long page_size = sysconf(_SC_PAGESIZE);  // in case x86-64 is configured to use 2MB pages
0545     long memory_usage = resident * page_size  / 1024 / 1024 ;    
0546 #endif
0547   
0548     
0549     std::cout << "Working on slice " << i + 1 << std::endl;
0550     std::cout << "Current memory usage: " << memory_usage << " MB" << std::endl;
0551 
0552   }
0553   // ---------------------------------------------------------------------------
0554 
0555   std::unique_ptr<HepMC3::GenEvent> mergeSlice(int i) {
0556     auto hepSlice = std::make_unique<HepMC3::GenEvent>(HepMC3::Units::GEV, HepMC3::Units::MM);
0557     
0558     addFreqEvents(signalFile, sigAdapter, sigFreq, hepSlice, signalStatus, true);
0559     
0560     for (const auto& freqBgs : freqAdapters) {
0561       auto fileName=freqBgs.first;
0562       addFreqEvents(fileName, freqAdapters[fileName], freqs[fileName], hepSlice, baseStatuses[fileName], false);
0563     }
0564     
0565     for (const auto& fileName : weightDict) {
0566       addWeightedEvents(fileName.first, hepSlice, baseStatuses[fileName.first]);
0567     }
0568 
0569     return hepSlice;
0570   };
0571 
0572   // ---------------------------------------------------------------------------
0573 
0574   void addFreqEvents(std::string fileName, std::shared_ptr<HepMC3::Reader>& adapter, const double freq,
0575              std::unique_ptr<HepMC3::GenEvent>& hepSlice, int baseStatus = 0, bool signal = false) {
0576 
0577     // First, create a timeline
0578     // Signals can be different
0579     std::vector<double> timeline;
0580 
0581     std::uniform_real_distribution<> uni(0, intWindow);
0582     if (freq == 0){
0583       if (!signal) {
0584         std::cerr << "frequency can't be 0 for background files" << std::endl;
0585         exit(EXIT_FAILURE);
0586       }
0587       // exactly one signal event, at an arbitrary point
0588       timeline.push_back(uni(rng));
0589     } else {
0590       // Generate poisson-distributed times to place events
0591       timeline = poissonTimes(freq, intWindow);
0592     }
0593     
0594     if ( verbose) std::cout << "Placing " << timeline.size() << " events from " << fileName << std::endl;
0595 
0596     if (timeline.empty()) return;
0597     long particleCount = 0;
0598 
0599     // Insert events at all specified locations
0600     for (double time : timeline) {
0601       if(adapter->failed()) {
0602       try{
0603         if (signal) { // Exhausted signal events
0604             throw std::ifstream::failure("EOF");
0605           } else { // background file reached its end, reset to the start
0606             std::cout << "Cycling back to the start of " << fileName << std::endl;
0607             adapter->close();
0608             adapter = HepMC3::deduce_reader(fileName);
0609           }
0610         } catch (std::ifstream::failure& e) {
0611             continue; // just need to suppress the error
0612         }
0613       }
0614 
0615       HepMC3::GenEvent inevt;
0616       adapter->read_event(inevt);
0617       if (signal && (signalFreq == 0.0)){
0618         hepSlice->weights() = inevt.weights();
0619       }
0620 
0621       if (squashTime) time = 0;
0622       particleCount += insertHepmcEvent( inevt, hepSlice, time, baseStatus, signal);
0623     }
0624 
0625     infoDict[fileName].eventCount += timeline.size();
0626     infoDict[fileName].particleCount += particleCount;
0627 
0628 
0629     return;
0630   }
0631 
0632   // ---------------------------------------------------------------------------
0633 
0634   void addWeightedEvents(std::string fileName, std::unique_ptr<HepMC3::GenEvent>& hepSlice, int baseStatus=0, bool signal = false) {
0635     auto& [events, weightedDist, avgRate ] = weightDict[fileName];
0636 
0637     // How many events? Assume Poisson distribution
0638     int nEvents;
0639     std::poisson_distribution<> d( intWindow * avgRate );
0640 
0641     // Small SR files may not have enough photons (example or test files). Could use them all or reroll
0642     // Choosing the latter, neither is physical
0643     while (true) {
0644       nEvents = d(rng);
0645       if (nEvents > events.size()) {
0646           std::cout << "WARNING: Trying to place " << nEvents << " events from " << fileName
0647                   << " but the file doesn't have enough. Rerolling, but this is not physical." << std::endl;
0648         continue;
0649       }
0650       break;
0651     }
0652 
0653     if (verbose) std::cout << "Placing " << nEvents << " events from " << fileName << std::endl;
0654     
0655     // Get randomized event indices
0656     // Note: Could change to drawing without replacing ( if ( not in toPLace) ...) , not worth the effort
0657     std::vector<HepMC3::GenEvent> toPlace(nEvents);
0658     for ( auto& e : toPlace ){
0659       auto i = static_cast<int> (weightedDist(rng));
0660       e = events.at(i);
0661     }
0662     
0663     // Place at random times
0664     std::vector<double> timeline;
0665     std::uniform_real_distribution<> uni(0, intWindow);
0666     long particleCount = 0;
0667     if (!squashTime) {
0668       for ( auto& e : toPlace ){
0669           double time = squashTime ? 0 : uni(rng);
0670         particleCount += insertHepmcEvent( e, hepSlice, time, baseStatus, signal);
0671       }
0672     }
0673 
0674     infoDict[fileName].eventCount += nEvents;
0675     infoDict[fileName].particleCount += particleCount;
0676 
0677     return;
0678 }
0679 
0680   // ---------------------------------------------------------------------------
0681   long insertHepmcEvent( const HepMC3::GenEvent& inevt,
0682              std::unique_ptr<HepMC3::GenEvent>& hepSlice, double time=0, int baseStatus=0, bool signal = false) {
0683     // Unit conversion
0684     double timeHepmc = c_light * time;
0685     
0686     std::vector<HepMC3::GenParticlePtr> particles;
0687     std::vector<HepMC3::GenVertexPtr> vertices;
0688 
0689     // Stores the vertices of the event inside a vertex container. These vertices are in increasing order
0690     // so we can index them with [abs(vertex_id)-1]
0691     for (auto& vertex : inevt.vertices()) {
0692       HepMC3::FourVector position = vertex->position();
0693       position.set_t(position.t() + timeHepmc);
0694       auto v1 = std::make_shared<HepMC3::GenVertex>(position);
0695       vertices.push_back(v1);
0696     }
0697       
0698     // copies the particles and attaches them to their corresponding vertices
0699     long finalParticleCount = 0;
0700     for (auto& particle : inevt.particles()) {
0701       HepMC3::FourVector momentum = particle->momentum();
0702       int status = particle->status();
0703       if (status == 1 ) finalParticleCount++;
0704       int pid = particle->pid();
0705       status += baseStatus;
0706       auto p1 = std::make_shared<HepMC3::GenParticle> (momentum, pid, status);
0707       p1->set_generated_mass(particle->generated_mass());
0708       particles.push_back(p1);
0709       // since the beam particles do not have a production vertex they cannot be attached to a production vertex
0710       if (particle->production_vertex()->id() < 0) {
0711           int production_vertex = particle->production_vertex()->id();
0712           vertices[abs(production_vertex) - 1]->add_particle_out(p1);
0713           hepSlice->add_particle(p1);
0714       }
0715     
0716       // Adds particles with an end vertex to their end vertices
0717       if (particle->end_vertex()) {
0718           int end_vertex = particle->end_vertex()->id();
0719           vertices.at(abs(end_vertex) - 1)->add_particle_in(p1);    
0720       }
0721     }
0722 
0723     // Adds the vertices with the attached particles to the event
0724     for (auto& vertex : vertices) {
0725       hepSlice->add_vertex(vertex);
0726     }
0727     
0728     return finalParticleCount;
0729   }
0730 
0731   // ---------------------------------------------------------------------------
0732 
0733   std::vector<double> poissonTimes(double mu, double endTime) {
0734     std::exponential_distribution<> exp(mu);
0735     
0736     double t = 0;
0737     std::vector<double> ret;
0738     while (true) {
0739       double delt = exp(rng)*1e6;
0740       // cout << delt <<endl;
0741       t += delt;
0742       if (t >= endTime) {
0743     break;
0744       }
0745       ret.push_back(t);
0746     }
0747     return ret;
0748 }
0749   // ---------------------------------------------------------------------------
0750 
0751   
0752   // private:
0753   std::mt19937 rng;
0754   string signalFile;
0755   double signalFreq;
0756   int signalSkip;
0757   int signalStatus;
0758   std::vector<BackgroundConfig> backgroundFiles;  
0759   string outputFile;
0760   string outputFileName;
0761   bool rootFormat;
0762   double intWindow;
0763   int nSlices; // should be long, but argparse cannot read that
0764   bool squashTime;
0765   int rngSeed;  // should be unsigned, but argparse cannot read that
0766   bool verbose;
0767   
0768   const double c_light = 299.792458; // speed of light = 299.792458 mm/ns to get mm  
0769 };
0770 
0771 // =============================================================
0772 int main(int argc, char* argv[]) {
0773 
0774   auto t0 = std::chrono::high_resolution_clock::now();
0775   // Create an instance of SignalBackgroundMerger
0776   SignalBackgroundMerger sbm (argc, argv);
0777 
0778 
0779   sbm.merge();
0780 
0781   std::cout << "\n==================================================================\n";
0782   std::cout << "Overall running time: " << std::round(std::chrono::duration<double, std::chrono::minutes::period>(std::chrono::high_resolution_clock::now() - t0).count()) << " min" << std::endl;
0783   
0784   return EXIT_SUCCESS;
0785 }