TICL/plugins/SimTrackstersProducer.cc

0001 // Author: Felice Pantaleo, Leonardo Cristella - felice.pantaleo@cern.ch, leonardo.cristella@cern.ch
0002 // Date: 09/2021
0003
0004 // user include files
0005
0006 #include "FWCore/Framework/interface/ESHandle.h"
0007 #include "FWCore/Framework/interface/Event.h"
0008 #include "FWCore/Framework/interface/Frameworkfwd.h"
0009 #include "FWCore/Framework/interface/MakerMacros.h"
0010 #include "FWCore/Framework/interface/stream/EDProducer.h"
0011 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0012 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0013 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0014
0015 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0016 #include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
0017
0018 #include "DataFormats/HGCalReco/interface/Trackster.h"
0019
0020 #include "DataFormats/Common/interface/ValueMap.h"
0021 #include "SimDataFormats/Associations/interface/LayerClusterToSimClusterAssociator.h"
0022 #include "SimDataFormats/Associations/interface/LayerClusterToCaloParticleAssociator.h"
0023
0024 #include "SimDataFormats/CaloAnalysis/interface/CaloParticle.h"
0025 #include "SimDataFormats/CaloAnalysis/interface/SimCluster.h"
0026 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
0027
0028 #include "RecoHGCal/TICL/interface/commons.h"
0029
0030 #include "TrackstersPCA.h"
0031 #include <vector>
0032 #include <map>
0033 #include <iterator>
0034 #include <algorithm>
0035
0036 using namespace ticl;
0037
0038 class SimTrackstersProducer : public edm::stream::EDProducer<> {
0039 public:
0040   explicit SimTrackstersProducer(const edm::ParameterSet&);
0041   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0042
0043   void produce(edm::Event&, const edm::EventSetup&) override;
0044   void addTrackster(const int& index,
0045                     const std::vector<std::pair<edm::Ref<reco::CaloClusterCollection>, std::pair<float, float>>>& lcVec,
0046                     const std::vector<float>& inputClusterMask,
0047                     const float& fractionCut_,
0048                     const float& energy,
0049                     const int& pdgId,
0050                     const int& charge,
0051                     const edm::ProductID& seed,
0052                     const Trackster::IterationIndex iter,
0053                     std::vector<float>& output_mask,
0054                     std::vector<Trackster>& result);
0055
0056 private:
0057   std::string detector_;
0058   const bool doNose_ = false;
0059   const edm::EDGetTokenT<std::vector<reco::CaloCluster>> clusters_token_;
0060   const edm::EDGetTokenT<edm::ValueMap<std::pair<float, float>>> clustersTime_token_;
0061   const edm::EDGetTokenT<std::vector<float>> filtered_layerclusters_mask_token_;
0062
0063   const edm::EDGetTokenT<std::vector<SimCluster>> simclusters_token_;
0064   const edm::EDGetTokenT<std::vector<CaloParticle>> caloparticles_token_;
0065
0066   const edm::EDGetTokenT<hgcal::SimToRecoCollectionWithSimClusters> associatorMapSimClusterToReco_token_;
0067   const edm::EDGetTokenT<hgcal::SimToRecoCollection> associatorMapCaloParticleToReco_token_;
0068   const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geom_token_;
0069   hgcal::RecHitTools rhtools_;
0070   const double fractionCut_;
0071 };
0072 DEFINE_FWK_MODULE(SimTrackstersProducer);
0073
0074 SimTrackstersProducer::SimTrackstersProducer(const edm::ParameterSet& ps)
0075     : detector_(ps.getParameter<std::string>("detector")),
0076       doNose_(detector_ == "HFNose"),
0077       clusters_token_(consumes(ps.getParameter<edm::InputTag>("layer_clusters"))),
0078       clustersTime_token_(consumes(ps.getParameter<edm::InputTag>("time_layerclusters"))),
0079       filtered_layerclusters_mask_token_(consumes(ps.getParameter<edm::InputTag>("filtered_mask"))),
0080       simclusters_token_(consumes(ps.getParameter<edm::InputTag>("simclusters"))),
0081       caloparticles_token_(consumes(ps.getParameter<edm::InputTag>("caloparticles"))),
0082       associatorMapSimClusterToReco_token_(
0083           consumes(ps.getParameter<edm::InputTag>("layerClusterSimClusterAssociator"))),
0084       associatorMapCaloParticleToReco_token_(
0085           consumes(ps.getParameter<edm::InputTag>("layerClusterCaloParticleAssociator"))),
0086       geom_token_(esConsumes()),
0087       fractionCut_(ps.getParameter<double>("fractionCut")) {
0088   produces<TracksterCollection>();
0089   produces<std::vector<float>>();
0090   produces<TracksterCollection>("fromCPs");
0091   produces<std::vector<float>>("fromCPs");
0092   produces<std::map<uint, std::vector<uint>>>();
0093 }
0094
0095 void SimTrackstersProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0096   // hgcalMultiClusters
0097   edm::ParameterSetDescription desc;
0098   desc.add<std::string>("detector", "HGCAL");
0099   desc.add<edm::InputTag>("layer_clusters", edm::InputTag("hgcalLayerClusters"));
0100   desc.add<edm::InputTag>("time_layerclusters", edm::InputTag("hgcalLayerClusters", "timeLayerCluster"));
0101   desc.add<edm::InputTag>("filtered_mask", edm::InputTag("filteredLayerClustersSimTracksters", "ticlSimTracksters"));
0102   desc.add<edm::InputTag>("simclusters", edm::InputTag("mix", "MergedCaloTruth"));
0103   desc.add<edm::InputTag>("caloparticles", edm::InputTag("mix", "MergedCaloTruth"));
0104   desc.add<edm::InputTag>("layerClusterSimClusterAssociator",
0105                           edm::InputTag("layerClusterSimClusterAssociationProducer"));
0106   desc.add<edm::InputTag>("layerClusterCaloParticleAssociator",
0107                           edm::InputTag("layerClusterCaloParticleAssociationProducer"));
0108   desc.add<double>("fractionCut", 0.);
0109
0110   descriptions.addWithDefaultLabel(desc);
0111 }
0112
0113 void SimTrackstersProducer::addTrackster(
0114     const int& index,
0115     const std::vector<std::pair<edm::Ref<reco::CaloClusterCollection>, std::pair<float, float>>>& lcVec,
0116     const std::vector<float>& inputClusterMask,
0117     const float& fractionCut_,
0118     const float& energy,
0119     const int& pdgId,
0120     const int& charge,
0121     const edm::ProductID& seed,
0122     const Trackster::IterationIndex iter,
0123     std::vector<float>& output_mask,
0124     std::vector<Trackster>& result) {
0125   if (lcVec.empty())
0126     return;
0127
0128   Trackster tmpTrackster;
0129   tmpTrackster.zeroProbabilities();
0130   tmpTrackster.vertices().reserve(lcVec.size());
0131   tmpTrackster.vertex_multiplicity().reserve(lcVec.size());
0132   for (auto const& [lc, energyScorePair] : lcVec) {
0133     if (inputClusterMask[lc.index()] > 0) {
0134       double fraction = energyScorePair.first / lc->energy();
0135       if (fraction < fractionCut_)
0136         continue;
0137       tmpTrackster.vertices().push_back(lc.index());
0138       output_mask[lc.index()] -= fraction;
0139       tmpTrackster.vertex_multiplicity().push_back(1. / fraction);
0140     }
0141   }
0142
0143   tmpTrackster.setIdProbability(tracksterParticleTypeFromPdgId(pdgId, charge), 1.f);
0144   tmpTrackster.setRegressedEnergy(energy);
0145   tmpTrackster.setIteration(iter);
0146   tmpTrackster.setSeed(seed, index);
0147   result.emplace_back(tmpTrackster);
0148 }
0149
0150 void SimTrackstersProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
0151   auto result = std::make_unique<TracksterCollection>();
0152   auto output_mask = std::make_unique<std::vector<float>>();
0153   auto result_fromCP = std::make_unique<TracksterCollection>();
0154   auto output_mask_fromCP = std::make_unique<std::vector<float>>();
0155   auto cpToSc_SimTrackstersMap = std::make_unique<std::map<uint, std::vector<uint>>>();
0156   const auto& layerClusters = evt.get(clusters_token_);
0157   const auto& layerClustersTimes = evt.get(clustersTime_token_);
0158   const auto& inputClusterMask = evt.get(filtered_layerclusters_mask_token_);
0159   output_mask->resize(layerClusters.size(), 1.f);
0160   output_mask_fromCP->resize(layerClusters.size(), 1.f);
0161
0162   const auto& simclusters = evt.get(simclusters_token_);
0163   const auto& caloparticles = evt.get(caloparticles_token_);
0164
0165   const auto& simClustersToRecoColl = evt.get(associatorMapSimClusterToReco_token_);
0166   const auto& caloParticlesToRecoColl = evt.get(associatorMapCaloParticleToReco_token_);
0167
0168   const auto& geom = es.getData(geom_token_);
0169   rhtools_.setGeometry(geom);
0170   const auto num_simclusters = simclusters.size();
0171   result->reserve(num_simclusters);  // Conservative size, will call shrink_to_fit later
0172   const auto num_caloparticles = caloparticles.size();
0173   result_fromCP->reserve(num_caloparticles);
0174
0175   for (const auto& [key, lcVec] : caloParticlesToRecoColl) {
0176     auto const& cp = *(key);
0177     auto cpIndex = &cp - &caloparticles[0];
0178
0179     auto regr_energy = cp.energy();
0180     std::vector<uint> scSimTracksterIdx;
0181     scSimTracksterIdx.reserve(cp.simClusters().size());
0182
0183     // Create a Trackster from the object entering HGCal
0184     if (cp.g4Tracks()[0].crossedBoundary()) {
0185       regr_energy = cp.g4Tracks()[0].getMomentumAtBoundary().energy();
0186
0187       addTrackster(cpIndex,
0188                    lcVec,
0189                    inputClusterMask,
0190                    fractionCut_,
0191                    regr_energy,
0192                    cp.pdgId(),
0193                    cp.charge(),
0194                    key.id(),
0195                    ticl::Trackster::SIM,
0196                    *output_mask,
0197                    *result);
0198     } else {
0199       for (const auto& scRef : cp.simClusters()) {
0200         const auto& it = simClustersToRecoColl.find(scRef);
0201         if (it == simClustersToRecoColl.end())
0202           continue;
0203         const auto& lcVec = it->val;
0204         auto const& sc = *(scRef);
0205         auto const scIndex = &sc - &simclusters[0];
0206
0207         addTrackster(scIndex,
0208                      lcVec,
0209                      inputClusterMask,
0210                      fractionCut_,
0211                      sc.g4Tracks()[0].getMomentumAtBoundary().energy(),
0212                      sc.pdgId(),
0213                      sc.charge(),
0214                      scRef.id(),
0215                      ticl::Trackster::SIM,
0216                      *output_mask,
0217                      *result);
0218
0219         if (result->empty())
0220           continue;
0221         const auto index = result->size() - 1;
0222         if (std::find(scSimTracksterIdx.begin(), scSimTracksterIdx.end(), index) == scSimTracksterIdx.end()) {
0223           scSimTracksterIdx.emplace_back(index);
0224         }
0225       }
0226       scSimTracksterIdx.shrink_to_fit();
0227     }
0228
0229     // Create a Trackster from any CP
0230     addTrackster(cpIndex,
0231                  lcVec,
0232                  inputClusterMask,
0233                  fractionCut_,
0234                  regr_energy,
0235                  cp.pdgId(),
0236                  cp.charge(),
0237                  key.id(),
0238                  ticl::Trackster::SIM_CP,
0239                  *output_mask_fromCP,
0240                  *result_fromCP);
0241
0242     if (result_fromCP->empty())
0243       continue;
0244     const auto index = result_fromCP->size() - 1;
0245     if (cpToSc_SimTrackstersMap->find(index) == cpToSc_SimTrackstersMap->end()) {
0246       (*cpToSc_SimTrackstersMap)[index] = scSimTracksterIdx;
0247     }
0248   }
0249
0250   ticl::assignPCAtoTracksters(
0251       *result, layerClusters, layerClustersTimes, rhtools_.getPositionLayer(rhtools_.lastLayerEE(doNose_)).z());
0252   result->shrink_to_fit();
0253   ticl::assignPCAtoTracksters(
0254       *result_fromCP, layerClusters, layerClustersTimes, rhtools_.getPositionLayer(rhtools_.lastLayerEE(doNose_)).z());
0255   result_fromCP->shrink_to_fit();
0256
0257   evt.put(std::move(result));
0258   evt.put(std::move(output_mask));
0259   evt.put(std::move(result_fromCP), "fromCPs");
0260   evt.put(std::move(output_mask_fromCP), "fromCPs");
0261   evt.put(std::move(cpToSc_SimTrackstersMap));
0262 }