Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-10-16 05:06:39

0001 // Author: Felice Pantaleo, felice.pantaleo@cern.ch 08/2024
0002 
0003 #include "FWCore/Framework/interface/global/EDProducer.h"
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/MakerMacros.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0007 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0008 #include "FWCore/Utilities/interface/EDGetToken.h"
0009 #include "DataFormats/HGCalReco/interface/Trackster.h"
0010 #include "SimDataFormats/Associations/interface/TICLAssociationMap.h"
0011 #include "DataFormats/Provenance/interface/ProductID.h"
0012 #include "FWCore/Framework/interface/ESHandle.h"
0013 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0014 
0015 class AllTracksterToSimTracksterAssociatorsByLCsProducer : public edm::global::EDProducer<> {
0016 public:
0017   explicit AllTracksterToSimTracksterAssociatorsByLCsProducer(const edm::ParameterSet&);
0018   ~AllTracksterToSimTracksterAssociatorsByLCsProducer() override = default;
0019 
0020   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0021 
0022 private:
0023   void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
0024 
0025   std::vector<std::pair<std::string, edm::EDGetTokenT<std::vector<ticl::Trackster>>>> tracksterCollectionTokens_;
0026   std::vector<std::pair<std::string, edm::EDGetTokenT<std::vector<ticl::Trackster>>>> simTracksterCollectionTokens_;
0027   edm::EDGetTokenT<std::vector<reco::CaloCluster>> layerClustersToken_;
0028   std::vector<std::pair<
0029       std::string,
0030       edm::EDGetTokenT<
0031           ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>>>
0032       layerClusterToTracksterMapTokens_;
0033   std::vector<std::pair<
0034       std::string,
0035       edm::EDGetTokenT<
0036           ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>>>
0037       layerClusterToSimTracksterMapTokens_;
0038 };
0039 
0040 AllTracksterToSimTracksterAssociatorsByLCsProducer::AllTracksterToSimTracksterAssociatorsByLCsProducer(
0041     const edm::ParameterSet& pset)
0042     : layerClustersToken_(consumes<std::vector<reco::CaloCluster>>(pset.getParameter<edm::InputTag>("layerClusters"))) {
0043   const auto& tracksterCollections = pset.getParameter<std::vector<edm::InputTag>>("tracksterCollections");
0044   for (const auto& tag : tracksterCollections) {
0045     std::string label = tag.label();
0046     if (tag.instance() != "") {
0047       label += tag.instance();
0048     }
0049     tracksterCollectionTokens_.emplace_back(label, consumes<std::vector<ticl::Trackster>>(tag));
0050     layerClusterToTracksterMapTokens_.emplace_back(
0051         label,
0052         consumes<
0053             ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>(
0054             edm::InputTag("allLayerClusterToTracksterAssociations", label)));
0055   }
0056 
0057   const auto& simTracksterCollections = pset.getParameter<std::vector<edm::InputTag>>("simTracksterCollections");
0058   for (const auto& tag : simTracksterCollections) {
0059     std::string label = tag.label();
0060     if (tag.instance() != "") {
0061       label += tag.instance();
0062     }
0063     simTracksterCollectionTokens_.emplace_back(label, consumes<std::vector<ticl::Trackster>>(tag));
0064     layerClusterToSimTracksterMapTokens_.emplace_back(
0065         label,
0066         consumes<
0067             ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>(
0068             edm::InputTag("allLayerClusterToTracksterAssociations", label)));
0069   }
0070 
0071   // Produce separate association maps for each trackster-simTrackster combination
0072   for (const auto& tracksterToken : tracksterCollectionTokens_) {
0073     for (const auto& simTracksterToken : simTracksterCollectionTokens_) {
0074       std::string instanceLabel = tracksterToken.first + "To" + simTracksterToken.first;
0075       produces<ticl::AssociationMap<ticl::mapWithFractionAndScore,
0076                                     std::vector<ticl::Trackster>,
0077                                     std::vector<ticl::Trackster>>>(instanceLabel);
0078       std::string reverseInstanceLabel = simTracksterToken.first + "To" + tracksterToken.first;
0079       produces<ticl::AssociationMap<ticl::mapWithFractionAndScore,
0080                                     std::vector<ticl::Trackster>,
0081                                     std::vector<ticl::Trackster>>>(reverseInstanceLabel);
0082     }
0083   }
0084 }
0085 
0086 void AllTracksterToSimTracksterAssociatorsByLCsProducer::produce(edm::StreamID,
0087                                                                  edm::Event& iEvent,
0088                                                                  const edm::EventSetup&) const {
0089   using namespace edm;
0090 
0091   Handle<std::vector<reco::CaloCluster>> layerClustersHandle;
0092   iEvent.getByToken(layerClustersToken_, layerClustersHandle);
0093   const auto& layerClusters = *layerClustersHandle;
0094 
0095   for (const auto& tracksterToken : tracksterCollectionTokens_) {
0096     Handle<std::vector<ticl::Trackster>> recoTrackstersHandle;
0097     iEvent.getByToken(tracksterToken.second, recoTrackstersHandle);
0098     const auto& recoTracksters = *recoTrackstersHandle;
0099 
0100     // Retrieve the correct LayerClusterToTracksterMap for the current trackster collection
0101     Handle<ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>
0102         layerClusterToTracksterMapHandle;
0103     auto tracksterMapTokenIter =
0104         std::find_if(layerClusterToTracksterMapTokens_.begin(),
0105                      layerClusterToTracksterMapTokens_.end(),
0106                      [&tracksterToken](const auto& pair) { return pair.first == tracksterToken.first; });
0107     if (tracksterMapTokenIter != layerClusterToTracksterMapTokens_.end()) {
0108       iEvent.getByToken(tracksterMapTokenIter->second, layerClusterToTracksterMapHandle);
0109     }
0110     const auto& layerClusterToTracksterMap = *layerClusterToTracksterMapHandle;
0111 
0112     for (const auto& simTracksterToken : simTracksterCollectionTokens_) {
0113       Handle<std::vector<ticl::Trackster>> simTrackstersHandle;
0114       iEvent.getByToken(simTracksterToken.second, simTrackstersHandle);
0115       const auto& simTracksters = *simTrackstersHandle;
0116 
0117       // Retrieve the correct LayerClusterToSimTracksterMap for the current simTrackster collection
0118       Handle<ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>
0119           layerClusterToSimTracksterMapHandle;
0120       auto simTracksterMapTokenIter =
0121           std::find_if(layerClusterToSimTracksterMapTokens_.begin(),
0122                        layerClusterToSimTracksterMapTokens_.end(),
0123                        [&simTracksterToken](const auto& pair) { return pair.first == simTracksterToken.first; });
0124       if (simTracksterMapTokenIter != layerClusterToSimTracksterMapTokens_.end()) {
0125         iEvent.getByToken(simTracksterMapTokenIter->second, layerClusterToSimTracksterMapHandle);
0126       }
0127       const auto& layerClusterToSimTracksterMap = *layerClusterToSimTracksterMapHandle;
0128 
0129       // Create the association maps
0130       auto tracksterToSimTracksterMap = std::make_unique<
0131           ticl::AssociationMap<ticl::mapWithFractionAndScore, std::vector<ticl::Trackster>, std::vector<ticl::Trackster>>>(
0132           recoTrackstersHandle, simTrackstersHandle, iEvent);
0133       auto simTracksterToTracksterMap = std::make_unique<
0134           ticl::AssociationMap<ticl::mapWithFractionAndScore, std::vector<ticl::Trackster>, std::vector<ticl::Trackster>>>(
0135           simTrackstersHandle, recoTrackstersHandle, iEvent);
0136 
0137       for (unsigned int tracksterIndex = 0; tracksterIndex < recoTracksters.size(); ++tracksterIndex) {
0138         const auto& recoTrackster = recoTracksters[tracksterIndex];
0139         edm::Ref<std::vector<ticl::Trackster>> recoTracksterRef(recoTrackstersHandle, tracksterIndex);
0140         const auto& layerClustersIds = recoTrackster.vertices();
0141         float recoToSimScoresDenominator = 0.f;
0142         ticl::mapWithFraction layerClusterToAssociatedSimTracksterMap(layerClustersIds.size());
0143         std::vector<unsigned int> associatedSimTracksterIndices;
0144         for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0145           unsigned int layerClusterId = layerClustersIds[i];
0146           const auto& layerCluster = layerClusters[layerClusterId];
0147           float recoFraction = 1.f / recoTrackster.vertex_multiplicity(i);
0148           float squaredRecoFraction = recoFraction * recoFraction;
0149           float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0150           recoToSimScoresDenominator += squaredLayerClusterEnergy * squaredRecoFraction;
0151           const auto& simTracksterVec = layerClusterToSimTracksterMap[layerClusterId];
0152           for (const auto& [simTracksterIndex, simSharedEnergy] : simTracksterVec) {
0153             layerClusterToAssociatedSimTracksterMap[i].push_back({simTracksterIndex, simSharedEnergy});
0154             associatedSimTracksterIndices.push_back(simTracksterIndex);
0155           }
0156         }
0157 
0158         // Keep only unique associatedSimTracksterIndices
0159         std::sort(associatedSimTracksterIndices.begin(), associatedSimTracksterIndices.end());
0160         associatedSimTracksterIndices.erase(
0161             std::unique(associatedSimTracksterIndices.begin(), associatedSimTracksterIndices.end()),
0162             associatedSimTracksterIndices.end());
0163 
0164         // Add missing sim tracksters with 0 shared energy to layerClusterToAssociatedSimTracksterMap
0165         for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0166           unsigned int layerClusterId = layerClustersIds[i];
0167           const auto& simTracksterVec = layerClusterToSimTracksterMap[layerClusterId];
0168           for (unsigned int simTracksterIndex : associatedSimTracksterIndices) {
0169             if (std::find_if(simTracksterVec.begin(), simTracksterVec.end(), [simTracksterIndex](const auto& pair) {
0170                   return pair.first == simTracksterIndex;
0171                 }) == simTracksterVec.end()) {
0172               layerClusterToAssociatedSimTracksterMap[i].push_back({simTracksterIndex, 0.f});
0173             }
0174           }
0175         }
0176 
0177         const float invDenominator = 1.f / recoToSimScoresDenominator;
0178 
0179         for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0180           unsigned int layerClusterId = layerClustersIds[i];
0181           const auto& layerCluster = layerClusters[layerClusterId];
0182           float recoFraction = 1.f / recoTrackster.vertex_multiplicity(i);
0183           float squaredRecoFraction = recoFraction * recoFraction;
0184           float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0185           float recoSharedEnergy = layerCluster.energy() * recoFraction;
0186           float invLayerClusterEnergy = 1.f / layerCluster.energy();
0187           const auto& simTracksterVec = layerClusterToAssociatedSimTracksterMap[i];
0188           for (const auto& [simTracksterIndex, simSharedEnergy] : simTracksterVec) {
0189             edm::Ref<std::vector<ticl::Trackster>> simTracksterRef(simTrackstersHandle, simTracksterIndex);
0190             float sharedEnergy = std::min(simSharedEnergy, recoSharedEnergy);
0191             float simFraction = simSharedEnergy * invLayerClusterEnergy;
0192             float score = invDenominator *
0193                           std::min(squaredRecoFraction, (recoFraction - simFraction) * (recoFraction - simFraction)) *
0194                           squaredLayerClusterEnergy;
0195             tracksterToSimTracksterMap->insert(recoTracksterRef, simTracksterRef, sharedEnergy, score);
0196           }
0197         }
0198       }
0199 
0200       for (unsigned int tracksterIndex = 0; tracksterIndex < simTracksters.size(); ++tracksterIndex) {
0201         const auto& simTrackster = simTracksters[tracksterIndex];
0202         edm::Ref<std::vector<ticl::Trackster>> simTracksterRef(simTrackstersHandle, tracksterIndex);
0203         const auto& layerClustersIds = simTrackster.vertices();
0204         float simToRecoScoresDenominator = 0.f;
0205         ticl::mapWithFraction layerClusterToAssociatedTracksterMap(layerClustersIds.size());
0206         std::vector<unsigned int> associatedRecoTracksterIndices;
0207         for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0208           unsigned int layerClusterId = layerClustersIds[i];
0209           const auto& layerCluster = layerClusters[layerClusterId];
0210           float simFraction = 1.f / simTrackster.vertex_multiplicity(i);
0211           float squaredSimFraction = simFraction * simFraction;
0212           float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0213           simToRecoScoresDenominator += squaredLayerClusterEnergy * squaredSimFraction;
0214           const auto& recoTracksterVec = layerClusterToTracksterMap[layerClusterId];
0215           for (const auto& [recoTracksterIndex, recoSharedEnergy] : recoTracksterVec) {
0216             layerClusterToAssociatedTracksterMap[i].push_back({recoTracksterIndex, recoSharedEnergy});
0217             associatedRecoTracksterIndices.push_back(recoTracksterIndex);
0218           }
0219         }
0220         // keep only unique associatedRecoTracksterIndices
0221         std::sort(associatedRecoTracksterIndices.begin(), associatedRecoTracksterIndices.end());
0222         associatedRecoTracksterIndices.erase(
0223             std::unique(associatedRecoTracksterIndices.begin(), associatedRecoTracksterIndices.end()),
0224             associatedRecoTracksterIndices.end());
0225         // for each layer cluster, loop over associatedRecoTracksterIndices and add the missing reco tracksters with 0 shared energy to layerClusterToAssociatedTracksterMap
0226         for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0227           unsigned int layerClusterId = layerClustersIds[i];
0228           const auto& recoTracksterVec = layerClusterToTracksterMap[layerClusterId];
0229           for (unsigned int recoTracksterIndex : associatedRecoTracksterIndices) {
0230             if (std::find_if(recoTracksterVec.begin(), recoTracksterVec.end(), [recoTracksterIndex](const auto& pair) {
0231                   return pair.first == recoTracksterIndex;
0232                 }) == recoTracksterVec.end()) {
0233               layerClusterToAssociatedTracksterMap[i].push_back({recoTracksterIndex, 0.f});
0234             }
0235           }
0236         }
0237 
0238         const float invDenominator = 1.f / simToRecoScoresDenominator;
0239 
0240         for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0241           unsigned int layerClusterId = layerClustersIds[i];
0242           const auto& layerCluster = layerClusters[layerClusterId];
0243           float simFraction = 1.f / simTrackster.vertex_multiplicity(i);
0244           float squaredSimFraction = simFraction * simFraction;
0245           float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0246           float simSharedEnergy = layerCluster.energy() * simFraction;
0247           float invLayerClusterEnergy = 1.f / layerCluster.energy();
0248           const auto& recoTracksterVec = layerClusterToAssociatedTracksterMap[i];
0249           for (const auto& [recoTracksterIndex, recoSharedEnergy] : recoTracksterVec) {
0250             edm::Ref<std::vector<ticl::Trackster>> recoTracksterRef(recoTrackstersHandle, recoTracksterIndex);
0251             float sharedEnergy = std::min(recoSharedEnergy, simSharedEnergy);
0252             float recoFraction = recoSharedEnergy * invLayerClusterEnergy;
0253             float score = invDenominator *
0254                           std::min(squaredSimFraction, (simFraction - recoFraction) * (simFraction - recoFraction)) *
0255                           squaredLayerClusterEnergy;
0256             simTracksterToTracksterMap->insert(tracksterIndex, recoTracksterIndex, sharedEnergy, score);
0257           }
0258         }
0259       }
0260       tracksterToSimTracksterMap->sort(true);
0261       simTracksterToTracksterMap->sort(true);
0262 
0263       // After populating the maps, store them in the event
0264       iEvent.put(std::move(tracksterToSimTracksterMap), tracksterToken.first + "To" + simTracksterToken.first);
0265       iEvent.put(std::move(simTracksterToTracksterMap), simTracksterToken.first + "To" + tracksterToken.first);
0266     }
0267   }
0268 }
0269 
0270 void AllTracksterToSimTracksterAssociatorsByLCsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0271   edm::ParameterSetDescription desc;
0272   desc.add<std::vector<edm::InputTag>>(
0273       "tracksterCollections", {edm::InputTag("ticlTrackstersCLUE3DHigh"), edm::InputTag("ticlTrackstersLinks")});
0274   desc.add<std::vector<edm::InputTag>>(
0275       "simTracksterCollections", {edm::InputTag("ticlSimTracksters"), edm::InputTag("ticlSimTracksters", "fromCPs")});
0276   desc.add<edm::InputTag>("layerClusters", edm::InputTag("hgcalMergeLayerClusters"));
0277   descriptions.add("AllTracksterToSimTracksterAssociatorsByLCsProducer", desc);
0278 }
0279 
0280 // Define this as a plug-in
0281 DEFINE_FWK_MODULE(AllTracksterToSimTracksterAssociatorsByLCsProducer);