utils/src/event_merge.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Wouter Deconinck
0003
0004 #include <regex>
0005 #include <tuple>
0006 #include <utility>
0007
0008 #include <CLI/CLI.hpp>
0009
0010 #include "podio/EventStore.h"
0011 #include "podio/ROOTReader.h"
0012 #include "podio/ROOTWriter.h"
0013
0014 #include "edm4hep/SimCalorimeterHitCollection.h"
0015 #include "edm4hep/SimTrackerHitCollection.h"
0016
0017 int main(int argc, char **argv) {
0018   // setup CLI options
0019   CLI::App app;
0020
0021   // debug
0022   bool debug{false};
0023   app.add_flag("--debug,-d", debug, "Enable debug output");
0024
0025   // verbose
0026   bool verbose{false};
0027   app.add_flag("--verbose,-v", verbose, "Enable verbose output");
0028
0029   // signal input file
0030   std::string file_sig{""};
0031   app.add_option("--signal,-s", file_sig, "Signal file");
0032
0033   // background input files
0034   std::vector<std::tuple<std::string, unsigned int>> files_bkg{};
0035   app.add_option("--background,-b", files_bkg, "Background files");
0036
0037   // output file
0038   std::string file_out{""};
0039   app.add_option("--output,-o", file_out, "Output file");
0040
0041   // collection include regex
0042   std::vector<std::string> include_regex{{".*"}};
0043   app.add_option("--include,-i", include_regex, "Collection inclusion regex");
0044
0045   // collection exclude regex
0046   std::vector<std::string> exclude_regex{};
0047   app.add_option("--exclude,-e", exclude_regex, "Collection exclusion regex");
0048
0049   // signal input file
0050   unsigned int numberOfEvents{0};
0051   app.add_option("--numberOfEvents,-n", numberOfEvents, "Number of events");
0052
0053   CLI11_PARSE(app, argc, argv);
0054
0055   if (debug) verbose = true;
0056
0057   // input reader
0058   auto reader_sig = podio::ROOTReader();
0059   auto store_sig = podio::EventStore();
0060   reader_sig.openFile(file_sig);
0061   store_sig.setReader(&reader_sig);
0062
0063   // background readers
0064   std::vector<std::tuple<podio::ROOTReader*, podio::EventStore*, unsigned int>> readers_stores_counts_bkg;
0065   for (const auto& [file_bkg, count_bkg]: files_bkg) {
0066     readers_stores_counts_bkg.emplace_back(new podio::ROOTReader(), new podio::EventStore(), count_bkg);
0067     std::get<0>(readers_stores_counts_bkg.back())->openFile(file_bkg);
0068     std::get<1>(readers_stores_counts_bkg.back())->setReader(std::get<0>(readers_stores_counts_bkg.back()));
0069   }
0070
0071   // output writer
0072   auto store_out = podio::EventStore();
0073   auto writer_out = podio::ROOTWriter(file_out, &store_out);
0074
0075   // hit collections
0076   auto coll_id_table_sig = store_sig.getCollectionIDTable();
0077   std::vector<std::pair<size_t, std::string>> calorimeter_collection_names;
0078   std::vector<std::pair<size_t, std::string>> tracker_collection_names;
0079   for (const auto& name: coll_id_table_sig->names()) {
0080     // check for inclusion
0081     for (const auto& include_re: include_regex) {
0082       if (std::regex_match(name, std::regex(include_re))) {
0083         // check for exclusion
0084         bool exclude = false;
0085         for (const auto& exclude_re: exclude_regex) {
0086           exclude = exclude || std::regex_match(name, std::regex(exclude_re));
0087         }
0088         // check for type by getting FIXME does not work
0089         auto& hits_sig = store_sig.get<edm4hep::SimTrackerHitCollection>(name);
0090         if (! store_sig.isValid()) {
0091           std::cout << "Collection " << name << " invalid as tracker collection" << std::endl;
0092           exclude = false;
0093         }
0094         //
0095         if (exclude == false) {
0096           std::cout << "Collection " << name << " included" << std::endl;
0097           tracker_collection_names.push_back(std::make_pair(tracker_collection_names.size(), name));
0098         }
0099       }
0100     }
0101   }
0102
0103   // create output calorimeter collections
0104   std::vector<edm4hep::SimCalorimeterHitCollection*> calorimeter_collections_out;
0105   for (const auto& [i, calorimeter_collection_name]: calorimeter_collection_names) {
0106     calorimeter_collections_out.push_back(&(store_out.create<edm4hep::SimCalorimeterHitCollection>(calorimeter_collection_name)));
0107     writer_out.registerForWrite(calorimeter_collection_name);
0108   }
0109
0110   // create output tracker collections
0111   std::vector<edm4hep::SimTrackerHitCollection*> tracker_collections_out;
0112   for (const auto& [i, tracker_collection_name]: tracker_collection_names) {
0113     tracker_collections_out.push_back(&(store_out.create<edm4hep::SimTrackerHitCollection>(tracker_collection_name)));
0114     writer_out.registerForWrite(tracker_collection_name);
0115   }
0116
0117   // loop over events
0118   unsigned n = std::min(numberOfEvents, reader_sig.getEntries());
0119   for (unsigned i = 0; i < n; ++i) {
0120     if (verbose) std::cout << "reading signal event " << i << "/" << n << std::endl;
0121
0122     // loop over signal collections
0123     for (const auto& [i, tracker_collection_name]: tracker_collection_names) {
0124       // clone signal hits
0125       auto& hits_sig = store_sig.get<edm4hep::SimTrackerHitCollection>(tracker_collection_name);
0126       if (debug) std::cout << " sig " << tracker_collection_name << ": " << hits_sig.size() << std::endl;
0127       for (const auto& hit_sig: hits_sig) {
0128         tracker_collections_out[i]->push_back(hit_sig.clone());
0129       }
0130     }
0131
0132     // loop over background collections
0133     for (const auto& [reader_bkg, store_bkg, count_bkg]: readers_stores_counts_bkg) {
0134       // repeat requested times
0135       for (auto count = 0; count != count_bkg; ++count) {
0136         // clone background hits (reader and store is pointer)
0137         for (const auto& [i, tracker_collection_name]: tracker_collection_names) {
0138           auto& hits_bkg = store_bkg->get<edm4hep::SimTrackerHitCollection>(tracker_collection_name);
0139           if (debug) std::cout << " bkg" << i << " " << tracker_collection_name << ": " << hits_bkg.size() << std::endl;
0140           for (const auto& hit_bkg: hits_bkg) {
0141             tracker_collections_out[i]->push_back(hit_bkg.clone());
0142           }
0143         }
0144         store_bkg->clear();
0145         reader_bkg->endOfEvent();
0146       }
0147     }
0148
0149     for (const auto& [i, tracker_collection_name]: tracker_collection_names) {
0150       if (debug) std::cout << " out " << tracker_collection_name << ": " << tracker_collections_out[i]->size() << std::endl;
0151     }
0152
0153     // clear input stores
0154     store_sig.clear();
0155     reader_sig.endOfEvent();
0156
0157     // write event
0158     writer_out.writeEvent();
0159     store_out.clearCollections();
0160   }
0161
0162   // close files
0163   reader_sig.closeFile();
0164   for (const auto& [reader_bkg, store_bkg, count_bkg]: readers_stores_counts_bkg) {
0165     reader_bkg->closeFile();
0166   }
0167   writer_out.finish();
0168
0169   return 0;
0170 }