TICL/plugins/TracksterLinkingbySuperClusteringMustache.cc

0001 #include <string>
0002 #include <memory>
0003 #include <algorithm>
0004
0005 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0007 #include "FWCore/ParameterSet/interface/allowedValues.h"
0008 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0009
0010 #include "DataFormats/HGCalReco/interface/TICLLayerTile.h"
0011 #include "DataFormats/HGCalReco/interface/Trackster.h"
0012 #include "RecoEcal/EgammaCoreTools/interface/Mustache.h"
0013 #include "RecoHGCal/TICL/plugins/TracksterLinkingbySuperClusteringMustache.h"
0014
0015 using namespace ticl;
0016
0017 TracksterLinkingbySuperClusteringMustache::TracksterLinkingbySuperClusteringMustache(
0018     const edm::ParameterSet& ps, edm::ConsumesCollector iC, cms::Ort::ONNXRuntime const* onnxRuntime)
0019     : TracksterLinkingAlgoBase(ps, iC, onnxRuntime),
0020       ecalMustacheSCParametersToken_(iC.esConsumes<EcalMustacheSCParameters, EcalMustacheSCParametersRcd>()),
0021       ecalSCDynamicDPhiParametersToken_(iC.esConsumes<EcalSCDynamicDPhiParameters, EcalSCDynamicDPhiParametersRcd>()),
0022       seedThresholdPt_(ps.getParameter<double>("seedThresholdPt")),
0023       candidateEnergyThreshold_(ps.getParameter<double>("candidateEnergyThreshold")),
0024       filterByTracksterPID_(ps.getParameter<bool>("filterByTracksterPID")),
0025       tracksterPIDCategoriesToFilter_(ps.getParameter<std::vector<int>>("tracksterPIDCategoriesToFilter")),
0026       PIDThreshold_(ps.getParameter<double>("PIDThreshold")) {}
0027
0028 void TracksterLinkingbySuperClusteringMustache::initialize(const HGCalDDDConstants* hgcons,
0029                                                            const hgcal::RecHitTools rhtools,
0030                                                            const edm::ESHandle<MagneticField> bfieldH,
0031                                                            const edm::ESHandle<Propagator> propH) {}
0032
0033 void TracksterLinkingbySuperClusteringMustache::setEvent(edm::Event& iEvent, edm::EventSetup const& iEventSetup) {
0034   mustacheSCParams_ = &iEventSetup.getData(ecalMustacheSCParametersToken_);
0035   scDynamicDPhiParams_ = &iEventSetup.getData(ecalSCDynamicDPhiParametersToken_);
0036 }
0037
0038 bool TracksterLinkingbySuperClusteringMustache::trackstersPassesPIDCut(const Trackster& tst) const {
0039   if (filterByTracksterPID_) {
0040     float probTotal = 0.0f;
0041     for (int cat : tracksterPIDCategoriesToFilter_) {
0042       probTotal += tst.id_probabilities(cat);
0043     }
0044     return probTotal >= PIDThreshold_;
0045   } else
0046     return true;
0047 }
0048
0049 /**
0050  * resultTracksters : superclusters as tracksters (ie merging of tracksters that have been superclustered together)
0051  * outputSuperclusters : same as linkedTracksterIdToInputTracksterId. Probably should use only one of the two.
0052  * linkedTracksterIdToInputTracksterId : maps indices from resultTracksters back into input tracksters.
0053  *    resultTracksters[i] has seed input.tracksters[linkedTracksterIdToInputTracksterId[i][0]], linked with tracksters input.tracksters[linkedTracksterIdToInputTracksterId[i][1..N]]
0054 */
0055 void TracksterLinkingbySuperClusteringMustache::linkTracksters(
0056     const Inputs& input,
0057     std::vector<Trackster>& resultTracksters,
0058     std::vector<std::vector<unsigned int>>& outputSuperclusters,
0059     std::vector<std::vector<unsigned int>>& linkedTracksterIdToInputTracksterId) {
0060   // For now we use all input tracksters for superclustering. At some point there might be a filter here for EM tracksters (electromagnetic identification with DNN ?)
0061   auto const& inputTracksters = input.tracksters;
0062   const unsigned int tracksterCount = inputTracksters.size();
0063
0064   /* Sorting tracksters by decreasing order of pT (out-of-place sort).
0065   inputTracksters[trackstersIndicesPt[0]], ..., inputTracksters[trackstersIndicesPt[N]] makes a list of tracksters sorted by decreasing pT
0066   Indices into this pT sorted collection will have the suffix _pt. Thus inputTracksters[index] and inputTracksters[trackstersIndicesPt[index_pt]] are correct
0067   */
0068   std::vector<unsigned int> trackstersIndicesPt(inputTracksters.size());
0069   std::iota(trackstersIndicesPt.begin(), trackstersIndicesPt.end(), 0);
0070   std::stable_sort(
0071       trackstersIndicesPt.begin(), trackstersIndicesPt.end(), [&inputTracksters](unsigned int i1, unsigned int i2) {
0072         return inputTracksters[i1].raw_pt() > inputTracksters[i2].raw_pt();
0073       });
0074
0075   std::vector<bool> tracksterMask_pt(tracksterCount, false);  // Mask for already superclustered tracksters
0076   // We also mask tracksters that don't pass the PID cut
0077   for (unsigned int ts_idx_pt = 0; ts_idx_pt < tracksterCount; ts_idx_pt++) {
0078     tracksterMask_pt[ts_idx_pt] = !trackstersPassesPIDCut(inputTracksters[trackstersIndicesPt[ts_idx_pt]]);
0079   }
0080
0081   for (unsigned int ts_seed_idx_pt = 0; ts_seed_idx_pt < tracksterCount; ts_seed_idx_pt++) {
0082     Trackster const& ts_seed = inputTracksters[trackstersIndicesPt[ts_seed_idx_pt]];
0083     if (ts_seed.raw_pt() <= seedThresholdPt_)
0084       break;  // Look only at seed tracksters passing threshold, take advantage of pt sorting for fast exit
0085     if (tracksterMask_pt[ts_seed_idx_pt])
0086       continue;  // Trackster does not pass PID cut
0087
0088     outputSuperclusters.emplace_back(std::initializer_list<unsigned int>{trackstersIndicesPt[ts_seed_idx_pt]});
0089     resultTracksters.emplace_back(inputTracksters[trackstersIndicesPt[ts_seed_idx_pt]]);
0090     linkedTracksterIdToInputTracksterId.emplace_back(
0091         std::initializer_list<unsigned int>{trackstersIndicesPt[ts_seed_idx_pt]});
0092
0093     for (unsigned int ts_cand_idx_pt = ts_seed_idx_pt + 1; ts_cand_idx_pt < tracksterCount; ts_cand_idx_pt++) {
0094       if (tracksterMask_pt[ts_cand_idx_pt])
0095         continue;  // Trackster is either already superclustered or did not pass PID cut
0096
0097       Trackster const& ts_cand = inputTracksters[trackstersIndicesPt[ts_cand_idx_pt]];
0098
0099       if (ts_cand.raw_energy() <= candidateEnergyThreshold_)
0100         continue;
0101
0102       const bool passes_dphi = reco::MustacheKernel::inDynamicDPhiWindow(scDynamicDPhiParams_,
0103                                                                          ts_seed.barycenter().eta(),
0104                                                                          ts_seed.barycenter().phi(),
0105                                                                          ts_cand.raw_energy(),
0106                                                                          ts_cand.barycenter().eta(),
0107                                                                          ts_cand.barycenter().phi());
0108
0109       if (passes_dphi && reco::MustacheKernel::inMustache(mustacheSCParams_,
0110                                                           ts_seed.barycenter().eta(),
0111                                                           ts_seed.barycenter().phi(),
0112                                                           ts_cand.raw_energy(),
0113                                                           ts_cand.barycenter().eta(),
0114                                                           ts_cand.barycenter().phi())) {
0115         outputSuperclusters.back().push_back(trackstersIndicesPt[ts_cand_idx_pt]);
0116         resultTracksters.back().mergeTracksters(ts_cand);
0117         linkedTracksterIdToInputTracksterId.back().push_back(trackstersIndicesPt[ts_cand_idx_pt]);
0118         tracksterMask_pt[ts_cand_idx_pt] = true;
0119       }
0120     }
0121   }
0122 }
0123
0124 void TracksterLinkingbySuperClusteringMustache::fillPSetDescription(edm::ParameterSetDescription& desc) {
0125   TracksterLinkingAlgoBase::fillPSetDescription(desc);  // adds algo_verbosity
0126   desc.add<double>("seedThresholdPt", 1.)
0127       ->setComment("Minimum transverse energy of trackster to be considered as seed of a supercluster");
0128   desc.add<double>("candidateEnergyThreshold", 0.15)
0129       ->setComment("Minimum energy of trackster to be considered as candidate for superclustering");
0130   desc.add<bool>("filterByTracksterPID", true)->setComment("Filter tracksters before superclustering by PID score");
0131   desc.add<std::vector<int>>(
0132           "tracksterPIDCategoriesToFilter",
0133           {static_cast<int>(Trackster::ParticleType::photon), static_cast<int>(Trackster::ParticleType::electron)})
0134       ->setComment("List of PID particle types (ticl::Trackster::ParticleType enum) to consider for PID filtering");
0135   desc.add<double>("PIDThreshold", 0.8)->setComment("PID score threshold");
0136 }