HGCalAssociatorProducers/plugins/TracksterToSimTracksterAssociatorProducer.cc

0001 // Author: Felice Pantaleo, felice.pantaleo@cern.ch 06/2024
0002
0003 #include "TracksterToSimTracksterAssociatorProducer.h"
0004 #include "SimDataFormats/Associations/interface/TICLAssociationMap.h"
0005
0006 TracksterToSimTracksterAssociatorProducer::TracksterToSimTracksterAssociatorProducer(const edm::ParameterSet& pset)
0007     : recoTracksterCollectionToken_(
0008           consumes<std::vector<ticl::Trackster>>(pset.getParameter<edm::InputTag>("tracksters"))),
0009       simTracksterCollectionToken_(
0010           consumes<std::vector<ticl::Trackster>>(pset.getParameter<edm::InputTag>("simTracksters"))),
0011       layerClustersCollectionToken_(
0012           consumes<std::vector<reco::CaloCluster>>(pset.getParameter<edm::InputTag>("layerClusters"))),
0013       LayerClusterToTracksterMapToken_(
0014           consumes<
0015               ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>(
0016               pset.getParameter<edm::InputTag>("tracksterMap"))),
0017       LayerClusterToSimTracksterMapToken_(
0018           consumes<
0019               ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>(
0020               pset.getParameter<edm::InputTag>("simTracksterMap"))) {
0021   produces<
0022       ticl::AssociationMap<ticl::mapWithFractionAndScore, std::vector<ticl::Trackster>, std::vector<ticl::Trackster>>>(
0023       "tracksterToSimTracksterMap");
0024   produces<
0025       ticl::AssociationMap<ticl::mapWithFractionAndScore, std::vector<ticl::Trackster>, std::vector<ticl::Trackster>>>(
0026       "simTracksterToTracksterMap");
0027 }
0028
0029 TracksterToSimTracksterAssociatorProducer::~TracksterToSimTracksterAssociatorProducer() {}
0030
0031 void TracksterToSimTracksterAssociatorProducer::produce(edm::StreamID,
0032                                                         edm::Event& iEvent,
0033                                                         const edm::EventSetup& iSetup) const {
0034   edm::Handle<std::vector<ticl::Trackster>> recoTrackstersHandle;
0035   iEvent.getByToken(recoTracksterCollectionToken_, recoTrackstersHandle);
0036   const auto& recoTracksters = *recoTrackstersHandle;
0037
0038   edm::Handle<std::vector<ticl::Trackster>> simTrackstersHandle;
0039   iEvent.getByToken(simTracksterCollectionToken_, simTrackstersHandle);
0040   const auto& simTracksters = *simTrackstersHandle;
0041
0042   edm::Handle<std::vector<reco::CaloCluster>> layerClustersHandle;
0043   iEvent.getByToken(layerClustersCollectionToken_, layerClustersHandle);
0044   const auto& layerClusters = *layerClustersHandle;
0045
0046   edm::Handle<ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>
0047       layerClusterToTracksterMapHandle;
0048   iEvent.getByToken(LayerClusterToTracksterMapToken_, layerClusterToTracksterMapHandle);
0049   const auto& layerClusterToTracksterMap = *layerClusterToTracksterMapHandle;
0050
0051   edm::Handle<ticl::AssociationMap<ticl::mapWithFraction, std::vector<reco::CaloCluster>, std::vector<ticl::Trackster>>>
0052       layerClusterToSimTracksterMapHandle;
0053   iEvent.getByToken(LayerClusterToSimTracksterMapToken_, layerClusterToSimTracksterMapHandle);
0054   const auto& layerClusterToSimTracksterMap = *layerClusterToSimTracksterMapHandle;
0055
0056   auto tracksterToSimTracksterMap = std::make_unique<
0057       ticl::AssociationMap<ticl::mapWithFractionAndScore, std::vector<ticl::Trackster>, std::vector<ticl::Trackster>>>(
0058       recoTrackstersHandle, simTrackstersHandle, iEvent);
0059   auto simTracksterToTracksterMap = std::make_unique<
0060       ticl::AssociationMap<ticl::mapWithFractionAndScore, std::vector<ticl::Trackster>, std::vector<ticl::Trackster>>>(
0061       simTrackstersHandle, recoTrackstersHandle, iEvent);
0062
0063   for (unsigned int tracksterIndex = 0; tracksterIndex < recoTracksters.size(); ++tracksterIndex) {
0064     const auto& recoTrackster = recoTracksters[tracksterIndex];
0065     edm::Ref<std::vector<ticl::Trackster>> recoTracksterRef(recoTrackstersHandle, tracksterIndex);
0066     const auto& layerClustersIds = recoTrackster.vertices();
0067     float recoToSimScoresDenominator = 0.f;
0068     ticl::mapWithFraction layerClusterToAssociatedSimTracksterMap(layerClustersIds.size());
0069     std::vector<unsigned int> associatedSimTracksterIndices;
0070     for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0071       unsigned int layerClusterId = layerClustersIds[i];
0072       const auto& layerCluster = layerClusters[layerClusterId];
0073       float recoFraction = 1.f / recoTrackster.vertex_multiplicity(i);
0074       float squaredRecoFraction = recoFraction * recoFraction;
0075       float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0076       recoToSimScoresDenominator += squaredLayerClusterEnergy * squaredRecoFraction;
0077       const auto& simTracksterVec = layerClusterToSimTracksterMap[layerClusterId];
0078       for (const auto& [simTracksterIndex, simSharedEnergy] : simTracksterVec) {
0079         layerClusterToAssociatedSimTracksterMap[i].push_back({simTracksterIndex, simSharedEnergy});
0080         associatedSimTracksterIndices.push_back(simTracksterIndex);
0081       }
0082     }
0083
0084     // Keep only unique associatedSimTracksterIndices
0085     std::sort(associatedSimTracksterIndices.begin(), associatedSimTracksterIndices.end());
0086     associatedSimTracksterIndices.erase(
0087         std::unique(associatedSimTracksterIndices.begin(), associatedSimTracksterIndices.end()),
0088         associatedSimTracksterIndices.end());
0089
0090     // Add missing sim tracksters with 0 shared energy to layerClusterToAssociatedSimTracksterMap
0091     for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0092       unsigned int layerClusterId = layerClustersIds[i];
0093       const auto& simTracksterVec = layerClusterToSimTracksterMap[layerClusterId];
0094       for (unsigned int simTracksterIndex : associatedSimTracksterIndices) {
0095         if (std::find_if(simTracksterVec.begin(), simTracksterVec.end(), [simTracksterIndex](const auto& pair) {
0096               return pair.first == simTracksterIndex;
0097             }) == simTracksterVec.end()) {
0098           layerClusterToAssociatedSimTracksterMap[i].push_back({simTracksterIndex, 0.f});
0099         }
0100       }
0101     }
0102
0103     const float invDenominator = 1.f / recoToSimScoresDenominator;
0104
0105     for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0106       unsigned int layerClusterId = layerClustersIds[i];
0107       const auto& layerCluster = layerClusters[layerClusterId];
0108       float recoFraction = 1.f / recoTrackster.vertex_multiplicity(i);
0109       float squaredRecoFraction = recoFraction * recoFraction;
0110       float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0111       float recoSharedEnergy = layerCluster.energy() * recoFraction;
0112       float invLayerClusterEnergy = 1.f / layerCluster.energy();
0113       const auto& simTracksterVec = layerClusterToAssociatedSimTracksterMap[i];
0114       for (const auto& [simTracksterIndex, simSharedEnergy] : simTracksterVec) {
0115         edm::Ref<std::vector<ticl::Trackster>> simTracksterRef(simTrackstersHandle, simTracksterIndex);
0116         float sharedEnergy = std::min(simSharedEnergy, recoSharedEnergy);
0117         float simFraction = simSharedEnergy * invLayerClusterEnergy;
0118         float score = invDenominator *
0119                       std::min(squaredRecoFraction, (recoFraction - simFraction) * (recoFraction - simFraction)) *
0120                       squaredLayerClusterEnergy;
0121         tracksterToSimTracksterMap->insert(recoTracksterRef, simTracksterRef, sharedEnergy, score);
0122       }
0123     }
0124   }
0125
0126   for (unsigned int tracksterIndex = 0; tracksterIndex < simTracksters.size(); ++tracksterIndex) {
0127     const auto& simTrackster = simTracksters[tracksterIndex];
0128     edm::Ref<std::vector<ticl::Trackster>> simTracksterRef(simTrackstersHandle, tracksterIndex);
0129     const auto& layerClustersIds = simTrackster.vertices();
0130     float simToRecoScoresDenominator = 0.f;
0131     ticl::mapWithFraction layerClusterToAssociatedTracksterMap(layerClustersIds.size());
0132     std::vector<unsigned int> associatedRecoTracksterIndices;
0133     for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0134       unsigned int layerClusterId = layerClustersIds[i];
0135       const auto& layerCluster = layerClusters[layerClusterId];
0136       float simFraction = 1.f / simTrackster.vertex_multiplicity(i);
0137       float squaredSimFraction = simFraction * simFraction;
0138       float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0139       simToRecoScoresDenominator += squaredLayerClusterEnergy * squaredSimFraction;
0140       const auto& recoTracksterVec = layerClusterToTracksterMap[layerClusterId];
0141       for (const auto& [recoTracksterIndex, recoSharedEnergy] : recoTracksterVec) {
0142         layerClusterToAssociatedTracksterMap[i].push_back({recoTracksterIndex, recoSharedEnergy});
0143         associatedRecoTracksterIndices.push_back(recoTracksterIndex);
0144       }
0145     }
0146     // keep only unique associatedRecoTracksterIndices
0147     std::sort(associatedRecoTracksterIndices.begin(), associatedRecoTracksterIndices.end());
0148     associatedRecoTracksterIndices.erase(
0149         std::unique(associatedRecoTracksterIndices.begin(), associatedRecoTracksterIndices.end()),
0150         associatedRecoTracksterIndices.end());
0151     // for each layer cluster, loop over associatedRecoTracksterIndices and add the missing reco tracksters with 0 shared energy to layerClusterToAssociatedTracksterMap
0152     for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0153       unsigned int layerClusterId = layerClustersIds[i];
0154       const auto& recoTracksterVec = layerClusterToTracksterMap[layerClusterId];
0155       for (unsigned int recoTracksterIndex : associatedRecoTracksterIndices) {
0156         if (std::find_if(recoTracksterVec.begin(), recoTracksterVec.end(), [recoTracksterIndex](const auto& pair) {
0157               return pair.first == recoTracksterIndex;
0158             }) == recoTracksterVec.end()) {
0159           layerClusterToAssociatedTracksterMap[i].push_back({recoTracksterIndex, 0.f});
0160         }
0161       }
0162     }
0163
0164     const float invDenominator = 1.f / simToRecoScoresDenominator;
0165
0166     for (unsigned int i = 0; i < layerClustersIds.size(); ++i) {
0167       unsigned int layerClusterId = layerClustersIds[i];
0168       const auto& layerCluster = layerClusters[layerClusterId];
0169       float simFraction = 1.f / simTrackster.vertex_multiplicity(i);
0170       float squaredSimFraction = simFraction * simFraction;
0171       float squaredLayerClusterEnergy = layerCluster.energy() * layerCluster.energy();
0172       float simSharedEnergy = layerCluster.energy() * simFraction;
0173       float invLayerClusterEnergy = 1.f / layerCluster.energy();
0174       const auto& recoTracksterVec = layerClusterToAssociatedTracksterMap[i];
0175       for (const auto& [recoTracksterIndex, recoSharedEnergy] : recoTracksterVec) {
0176         edm::Ref<std::vector<ticl::Trackster>> recoTracksterRef(recoTrackstersHandle, recoTracksterIndex);
0177         float sharedEnergy = std::min(recoSharedEnergy, simSharedEnergy);
0178         float recoFraction = recoSharedEnergy * invLayerClusterEnergy;
0179         float score = invDenominator *
0180                       std::min(squaredSimFraction, (simFraction - recoFraction) * (simFraction - recoFraction)) *
0181                       squaredLayerClusterEnergy;
0182         simTracksterToTracksterMap->insert(tracksterIndex, recoTracksterIndex, sharedEnergy, score);
0183       }
0184     }
0185   }
0186   tracksterToSimTracksterMap->sort(true);
0187   simTracksterToTracksterMap->sort(true);
0188   iEvent.put(std::move(tracksterToSimTracksterMap), "tracksterToSimTracksterMap");
0189   iEvent.put(std::move(simTracksterToTracksterMap), "simTracksterToTracksterMap");
0190 }
0191
0192 void TracksterToSimTracksterAssociatorProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0193   edm::ParameterSetDescription desc;
0194   desc.add<edm::InputTag>("tracksters", edm::InputTag("trackstersProducer"));
0195   desc.add<edm::InputTag>("simTracksters", edm::InputTag("simTrackstersProducer"));
0196   desc.add<edm::InputTag>("layerClusters", edm::InputTag("hgcalMergeLayerClusters"));
0197   desc.add<edm::InputTag>("tracksterMap", edm::InputTag("tracksterAssociatorProducer"));
0198   desc.add<edm::InputTag>("simTracksterMap", edm::InputTag("simTracksterAssociatorProducer"));
0199   descriptions.add("tracksterToSimTracksterAssociatorProducer", desc);
0200 }
0201
0202 DEFINE_FWK_MODULE(TracksterToSimTracksterAssociatorProducer);