File indexing completed on 2022-08-30 22:39:08
0001
0002
0003
0004
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
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);
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
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
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 }