Associations/interface/TICLAssociationMap.h

0001 #ifndef SimDataFormats_Associations_TICLAssociationMap_h
0002 #define SimDataFormats_Associations_TICLAssociationMap_h
0003
0004 #include <vector>
0005 #include <utility>
0006 #include <algorithm>
0007 #include <stdexcept>
0008 #include <type_traits>
0009 #include <iostream>
0010
0011 // CMSSW specific includes
0012 #include "DataFormats/Common/interface/Ref.h"
0013 #include "DataFormats/Common/interface/RefProd.h"
0014 #include "FWCore/Framework/interface/Event.h"
0015
0016 namespace ticl {
0017
0018   // Define the possible map types
0019   using mapWithFraction = std::vector<std::vector<std::pair<unsigned int, float>>>;
0020   using mapWithFractionAndScore = std::vector<std::vector<std::pair<unsigned int, std::pair<float, float>>>>;
0021   using oneToOneMapWithFraction = std::vector<std::pair<unsigned int, float>>;
0022   using oneToOneMapWithFractionAndScore = std::vector<std::pair<unsigned int, std::pair<float, float>>>;
0023
0024   template <typename MapType, typename Collection1 = void, typename Collection2 = void>
0025   class AssociationMap {
0026   private:
0027     // Type alias for conditionally including collectionRefProds
0028     using CollectionRefProdType =
0029         typename std::conditional_t<std::is_void_v<Collection1> || std::is_void_v<Collection2>,
0030                                     std::monostate,
0031                                     std::pair<edm::RefProd<Collection1>, edm::RefProd<Collection2>>>;
0032
0033   public:
0034     AssociationMap() : collectionRefProds() {}
0035     // Constructor for generic use
0036     template <typename C1 = Collection1,
0037               typename C2 = Collection2,
0038               typename std::enable_if_t<std::is_void_v<C1> && std::is_void_v<C2>, int> = 0>
0039     AssociationMap(const unsigned int size1 = 0) {
0040       map_.resize(size1);
0041     }
0042
0043     // Constructor for CMSSW-specific use
0044     template <typename C1 = Collection1,
0045               typename C2 = Collection2,
0046               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0047     AssociationMap(const edm::RefProd<C1>& id1, const edm::RefProd<C2>& id2, const edm::Event& event)
0048         : collectionRefProds(std::make_pair(id1, id2)) {
0049       resize(event);
0050     }
0051
0052     // Constructor for CMSSW-specific use
0053     template <typename C1 = Collection1,
0054               typename C2 = Collection2,
0055               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0056     AssociationMap(const edm::Handle<C1>& handle1, const edm::Handle<C2>& handle2, const edm::Event& event)
0057         : collectionRefProds(std::make_pair(edm::RefProd<C1>(handle1), edm::RefProd<C2>(handle2))) {
0058       resize(event);
0059     }
0060
0061     MapType& getMap() { return map_; }
0062
0063     const MapType& getMap() const { return map_; }
0064
0065     // CMSSW-specific method to get references
0066     template <typename C1 = Collection1,
0067               typename C2 = Collection2,
0068               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0069     edm::Ref<C1> getRefFirst(unsigned int index) const {
0070       return edm::Ref<C1>(collectionRefProds.first, index);
0071     }
0072
0073     template <typename C1 = Collection1,
0074               typename C2 = Collection2,
0075               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0076     edm::Ref<C2> getRefSecond(unsigned int index) const {
0077       return edm::Ref<C2>(collectionRefProds.second, index);
0078     }
0079
0080     // Method to get collection IDs for CMSSW-specific use
0081     template <typename C1 = Collection1,
0082               typename C2 = Collection2,
0083               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0084     std::pair<const edm::RefProd<C1>, const edm::RefProd<C2>> getCollectionIDs() const {
0085       return collectionRefProds;
0086     }
0087
0088     void insert(unsigned int index1, unsigned int index2, float fraction, float score = 0.0f) {
0089       if constexpr (std::is_same<MapType, mapWithFraction>::value) {
0090         if (index1 >= map_.size()) {
0091           map_.resize(index1 + 1);
0092         }
0093         auto& vec = map_[index1];
0094         auto it = std::find_if(vec.begin(), vec.end(), [&](const auto& pair) { return pair.first == index2; });
0095         if (it != vec.end()) {
0096           it->second += fraction;
0097         } else {
0098           vec.emplace_back(index2, fraction);
0099         }
0100       } else if constexpr (std::is_same<MapType, mapWithFractionAndScore>::value) {
0101         if (index1 >= map_.size()) {
0102           map_.resize(index1 + 1);
0103         }
0104         auto& vec = map_[index1];
0105         auto it = std::find_if(vec.begin(), vec.end(), [&](const auto& pair) { return pair.first == index2; });
0106         if (it != vec.end()) {
0107           it->second.first += fraction;
0108           it->second.second += score;
0109         } else {
0110           vec.emplace_back(index2, std::make_pair(fraction, score));
0111         }
0112       } else if constexpr (std::is_same<MapType, oneToOneMapWithFraction>::value) {
0113         auto it = std::find_if(map_.begin(), map_.end(), [&](const auto& pair) { return pair.first == index1; });
0114         if (it != map_.end()) {
0115           it->second += fraction;
0116         } else {
0117           map_.emplace_back(index1, fraction);
0118         }
0119       } else if constexpr (std::is_same<MapType, oneToOneMapWithFractionAndScore>::value) {
0120         auto it = std::find_if(map_.begin(), map_.end(), [&](const auto& pair) { return pair.first == index1; });
0121         if (it != map_.end()) {
0122           it->second.first += fraction;
0123           it->second.second += score;
0124         } else {
0125           map_.emplace_back(index1, std::make_pair(fraction, score));
0126         }
0127       }
0128     }
0129
0130     // Overload of insert for CMSSW-specific use
0131     template <typename C1 = Collection1,
0132               typename C2 = Collection2,
0133               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0134     void insert(const edm::Ref<C1>& ref1, const edm::Ref<C2>& ref2, float fraction, float score = 0.0f) {
0135       insert(ref1.key(), ref2.key(), fraction, score);
0136     }
0137
0138     void sort(bool byScore = false) {
0139       static_assert(!std::is_same_v<MapType, oneToOneMapWithFraction> &&
0140                         !std::is_same_v<MapType, oneToOneMapWithFractionAndScore>,
0141                     "Sort is not applicable for one-to-one maps");
0142
0143       if constexpr (std::is_same_v<MapType, mapWithFraction>) {
0144         for (auto& vec : map_) {
0145           std::sort(vec.begin(), vec.end(), [](const auto& a, const auto& b) { return a.second > b.second; });
0146         }
0147       } else if constexpr (std::is_same_v<MapType, mapWithFractionAndScore>) {
0148         for (auto& vec : map_) {
0149           if (byScore) {
0150             std::sort(
0151                 vec.begin(), vec.end(), [](const auto& a, const auto& b) { return a.second.second > b.second.second; });
0152           } else {
0153             std::sort(
0154                 vec.begin(), vec.end(), [](const auto& a, const auto& b) { return a.second.first > b.second.first; });
0155           }
0156         }
0157       }
0158     }
0159
0160     auto& operator[](unsigned int index1) { return map_[index1]; }
0161
0162     const auto& operator[](unsigned int index1) const { return map_[index1]; }
0163
0164     const auto& at(unsigned int index1) const {
0165       if (index1 >= map_.size()) {
0166         throw std::out_of_range("Index out of range");
0167       }
0168       return map_[index1];
0169     }
0170
0171     auto& at(unsigned int index1) {
0172       if (index1 >= map_.size()) {
0173         throw std::out_of_range("Index out of range");
0174       }
0175       return map_[index1];
0176     }
0177     // CMSSW-specific resize method
0178     template <typename C1 = Collection1,
0179               typename C2 = Collection2,
0180               typename std::enable_if_t<!std::is_void_v<C1> && !std::is_void_v<C2>, int> = 0>
0181     void resize(const edm::Event& event) {
0182       map_.resize(collectionRefProds.first->size());
0183     }
0184     // Constructor for generic use
0185     template <typename C1 = Collection1,
0186               typename C2 = Collection2,
0187               typename std::enable_if_t<std::is_void_v<C1> && std::is_void_v<C2>, int> = 0>
0188     void resize(const unsigned int size1) {
0189       map_.resize(size1);
0190     }
0191
0192     // Method to print the entire map
0193     void print(std::ostream& os) const {
0194       for (size_t i = 0; i < map_.size(); ++i) {
0195         os << "Index " << i << ":\n";
0196         for (const auto& pair : map_[i]) {
0197           os << "  (" << pair.first << ", ";
0198           if constexpr (std::is_same<MapType, mapWithFractionAndScore>::value ||
0199                         std::is_same<MapType, oneToOneMapWithFractionAndScore>::value) {
0200             os << pair.second.first << ", " << pair.second.second;
0201           } else {
0202             os << pair.second;
0203           }
0204           os << ")\n";
0205         }
0206       }
0207     }
0208
0209   private:
0210     // For CMSSW-specific use
0211     CollectionRefProdType collectionRefProds;
0212     MapType map_;  // Store the map directly
0213   };
0214
0215 }  // namespace ticl
0216
0217 #endif