File indexing completed on 2024-09-27 07:02:49
0001
0002
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
0019 CLI::App app;
0020
0021
0022 bool debug{false};
0023 app.add_flag("--debug,-d", debug, "Enable debug output");
0024
0025
0026 bool verbose{false};
0027 app.add_flag("--verbose,-v", verbose, "Enable verbose output");
0028
0029
0030 std::string file_sig{""};
0031 app.add_option("--signal,-s", file_sig, "Signal file");
0032
0033
0034 std::vector<std::tuple<std::string, unsigned int>> files_bkg{};
0035 app.add_option("--background,-b", files_bkg, "Background files");
0036
0037
0038 std::string file_out{""};
0039 app.add_option("--output,-o", file_out, "Output file");
0040
0041
0042 std::vector<std::string> include_regex{{".*"}};
0043 app.add_option("--include,-i", include_regex, "Collection inclusion regex");
0044
0045
0046 std::vector<std::string> exclude_regex{};
0047 app.add_option("--exclude,-e", exclude_regex, "Collection exclusion regex");
0048
0049
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
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
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
0072 auto store_out = podio::EventStore();
0073 auto writer_out = podio::ROOTWriter(file_out, &store_out);
0074
0075
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
0081 for (const auto& include_re: include_regex) {
0082 if (std::regex_match(name, std::regex(include_re))) {
0083
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
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
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
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
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
0123 for (const auto& [i, tracker_collection_name]: tracker_collection_names) {
0124
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
0133 for (const auto& [reader_bkg, store_bkg, count_bkg]: readers_stores_counts_bkg) {
0134
0135 for (auto count = 0; count != count_bkg; ++count) {
0136
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
0154 store_sig.clear();
0155 reader_sig.endOfEvent();
0156
0157
0158 writer_out.writeEvent();
0159 store_out.clearCollections();
0160 }
0161
0162
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 }