File indexing completed on 2025-03-14 23:36:37
0001
0002
0003
0004
0005 #include <vector>
0006
0007 #include "FWCore/Framework/interface/ESHandle.h"
0008 #include "FWCore/Framework/interface/Event.h"
0009 #include "FWCore/Framework/interface/Frameworkfwd.h"
0010 #include "FWCore/Framework/interface/MakerMacros.h"
0011 #include "FWCore/Framework/interface/stream/EDProducer.h"
0012 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0013 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0014 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0015 #include "FWCore/ParameterSet/interface/PluginDescription.h"
0016
0017 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0018 #include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
0019
0020 #include "DataFormats/HGCalReco/interface/Trackster.h"
0021 #include "DataFormats/HGCalReco/interface/TICLLayerTile.h"
0022 #include "DataFormats/HGCalReco/interface/TICLSeedingRegion.h"
0023 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0024
0025 #include "RecoHGCal/TICL/plugins/PatternRecognitionPluginFactory.h"
0026
0027 #include "RecoHGCal/TICL/interface/TracksterInferenceAlgoBase.h"
0028 #include "FWCore/Framework/interface/ConsumesCollector.h"
0029 #include "RecoHGCal/TICL/interface/TracksterInferenceAlgoFactory.h"
0030
0031 using namespace ticl;
0032 using namespace cms::Ort;
0033
0034 class TrackstersProducer : public edm::stream::EDProducer<> {
0035 public:
0036 explicit TrackstersProducer(const edm::ParameterSet&);
0037 ~TrackstersProducer() override {}
0038 static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0039
0040 void produce(edm::Event&, const edm::EventSetup&) override;
0041
0042
0043 static std::unique_ptr<TrackstersCache> initializeGlobalCache(const edm::ParameterSet&);
0044 static void globalEndJob(TrackstersCache*);
0045
0046 private:
0047 std::string detector_;
0048 bool doNose_;
0049 std::unique_ptr<PatternRecognitionAlgoBaseT<TICLLayerTiles>> myAlgo_;
0050 std::unique_ptr<PatternRecognitionAlgoBaseT<TICLLayerTilesHFNose>> myAlgoHFNose_;
0051 std::unique_ptr<TracksterInferenceAlgoBase> inferenceAlgo_;
0052 const edm::EDGetTokenT<std::vector<reco::CaloCluster>> clusters_token_;
0053 const edm::EDGetTokenT<std::vector<float>> filtered_layerclusters_mask_token_;
0054 const edm::EDGetTokenT<std::vector<float>> original_layerclusters_mask_token_;
0055 const edm::EDGetTokenT<edm::ValueMap<std::pair<float, float>>> clustersTime_token_;
0056 edm::EDGetTokenT<TICLLayerTiles> layer_clusters_tiles_token_;
0057 edm::EDGetTokenT<TICLLayerTilesHFNose> layer_clusters_tiles_hfnose_token_;
0058 const edm::EDGetTokenT<std::vector<TICLSeedingRegion>> seeding_regions_token_;
0059 const std::string itername_;
0060 ticl::Trackster::IterationIndex iterIndex_ = ticl::Trackster::IterationIndex(0);
0061 };
0062 DEFINE_FWK_MODULE(TrackstersProducer);
0063
0064 TrackstersProducer::TrackstersProducer(const edm::ParameterSet& ps)
0065 : detector_(ps.getParameter<std::string>("detector")),
0066 doNose_(detector_ == "HFNose"),
0067 clusters_token_(consumes<std::vector<reco::CaloCluster>>(ps.getParameter<edm::InputTag>("layer_clusters"))),
0068 filtered_layerclusters_mask_token_(consumes<std::vector<float>>(ps.getParameter<edm::InputTag>("filtered_mask"))),
0069 original_layerclusters_mask_token_(consumes<std::vector<float>>(ps.getParameter<edm::InputTag>("original_mask"))),
0070 clustersTime_token_(
0071 consumes<edm::ValueMap<std::pair<float, float>>>(ps.getParameter<edm::InputTag>("time_layerclusters"))),
0072 seeding_regions_token_(
0073 consumes<std::vector<TICLSeedingRegion>>(ps.getParameter<edm::InputTag>("seeding_regions"))),
0074 itername_(ps.getParameter<std::string>("itername")) {
0075 auto plugin = ps.getParameter<std::string>("patternRecognitionBy");
0076 auto pluginPSet = ps.getParameter<edm::ParameterSet>("pluginPatternRecognitionBy" + plugin);
0077 if (doNose_) {
0078 myAlgoHFNose_ = PatternRecognitionHFNoseFactory::get()->create(
0079 ps.getParameter<std::string>("patternRecognitionBy"), pluginPSet, consumesCollector());
0080 layer_clusters_tiles_hfnose_token_ =
0081 consumes<TICLLayerTilesHFNose>(ps.getParameter<edm::InputTag>("layer_clusters_hfnose_tiles"));
0082 } else {
0083 myAlgo_ = PatternRecognitionFactory::get()->create(
0084 ps.getParameter<std::string>("patternRecognitionBy"), pluginPSet, consumesCollector());
0085 layer_clusters_tiles_token_ = consumes<TICLLayerTiles>(ps.getParameter<edm::InputTag>("layer_clusters_tiles"));
0086 }
0087
0088
0089 std::string inferencePlugin = ps.getParameter<std::string>("inferenceAlgo");
0090 edm::ParameterSet inferencePSet = ps.getParameter<edm::ParameterSet>("pluginInferenceAlgo" + inferencePlugin);
0091 inferenceAlgo_ = std::unique_ptr<TracksterInferenceAlgoBase>(
0092 TracksterInferenceAlgoFactory::get()->create(inferencePlugin, inferencePSet));
0093
0094 if (itername_ == "TrkEM")
0095 iterIndex_ = ticl::Trackster::TRKEM;
0096 else if (itername_ == "EM")
0097 iterIndex_ = ticl::Trackster::EM;
0098 else if (itername_ == "Trk")
0099 iterIndex_ = ticl::Trackster::TRKHAD;
0100 else if (itername_ == "HAD")
0101 iterIndex_ = ticl::Trackster::HAD;
0102 else if (itername_ == "MIP")
0103 iterIndex_ = ticl::Trackster::MIP;
0104
0105 produces<std::vector<Trackster>>();
0106 produces<std::vector<float>>();
0107 }
0108
0109 void TrackstersProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0110
0111 edm::ParameterSetDescription desc;
0112 desc.add<std::string>("detector", "HGCAL");
0113 desc.add<edm::InputTag>("layer_clusters", edm::InputTag("hgcalMergeLayerClusters"));
0114 desc.add<edm::InputTag>("filtered_mask", edm::InputTag("filteredLayerClusters", "iterationLabelGoesHere"));
0115 desc.add<edm::InputTag>("original_mask", edm::InputTag("hgcalMergeLayerClusters", "InitialLayerClustersMask"));
0116 desc.add<edm::InputTag>("time_layerclusters", edm::InputTag("hgcalMergeLayerClusters", "timeLayerCluster"));
0117 desc.add<edm::InputTag>("layer_clusters_tiles", edm::InputTag("ticlLayerTileProducer"));
0118 desc.add<edm::InputTag>("layer_clusters_hfnose_tiles", edm::InputTag("ticlLayerTileHFNose"));
0119 desc.add<edm::InputTag>("seeding_regions", edm::InputTag("ticlSeedingRegionProducer"));
0120 desc.add<std::string>("patternRecognitionBy", "CA");
0121 desc.add<std::string>("itername", "unknown");
0122 desc.add<std::string>("inferenceAlgo", "TracksterInferenceByPFN");
0123
0124
0125 edm::ParameterSetDescription pluginDesc;
0126 pluginDesc.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "CA", true));
0127 desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByCA", pluginDesc);
0128
0129
0130 edm::ParameterSetDescription pluginDescClue3D;
0131 pluginDescClue3D.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "CLUE3D", true));
0132 desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByCLUE3D", pluginDescClue3D);
0133
0134
0135 edm::ParameterSetDescription pluginDescFastJet;
0136 pluginDescFastJet.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "FastJet", true));
0137 desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByFastJet", pluginDescFastJet);
0138
0139
0140 edm::ParameterSetDescription pluginDescRecovery;
0141 pluginDescRecovery.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "Recovery", true));
0142 desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByRecovery", pluginDescRecovery);
0143
0144
0145 edm::ParameterSetDescription inferenceDesc;
0146 inferenceDesc.addNode(edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByDNN", true));
0147 desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByDNN", inferenceDesc);
0148
0149 edm::ParameterSetDescription inferenceDescPFN;
0150 inferenceDescPFN.addNode(
0151 edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByPFN", true));
0152 desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByPFN", inferenceDescPFN);
0153
0154 edm::ParameterSetDescription inferenceDescANN;
0155 inferenceDescANN.addNode(
0156 edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByANN", true));
0157 desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByANN", inferenceDescANN);
0158
0159 edm::ParameterSetDescription inferenceDescCNNv4;
0160 inferenceDescCNNv4.addNode(
0161 edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByCNNv4", true));
0162 desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByCNNv4", inferenceDescCNNv4);
0163
0164 descriptions.add("trackstersProducer", desc);
0165 }
0166
0167 void TrackstersProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
0168 auto result = std::make_unique<std::vector<Trackster>>();
0169 auto initialResult = std::make_unique<std::vector<Trackster>>();
0170 auto output_mask = std::make_unique<std::vector<float>>();
0171
0172 const std::vector<float>& original_layerclusters_mask = evt.get(original_layerclusters_mask_token_);
0173 const auto& layerClusters = evt.get(clusters_token_);
0174 const auto& inputClusterMask = evt.get(filtered_layerclusters_mask_token_);
0175 const auto& layerClustersTimes = evt.get(clustersTime_token_);
0176 const auto& seeding_regions = evt.get(seeding_regions_token_);
0177
0178 std::unordered_map<int, std::vector<int>> seedToTrackstersAssociation;
0179
0180 if (!seeding_regions.empty()) {
0181 if (seeding_regions[0].index != -1) {
0182 auto numberOfSeedingRegions = seeding_regions.size();
0183 for (unsigned int i = 0; i < numberOfSeedingRegions; ++i) {
0184 seedToTrackstersAssociation.emplace(seeding_regions[i].index, 0);
0185 }
0186 }
0187
0188 if (doNose_) {
0189 const auto& layer_clusters_hfnose_tiles = evt.get(layer_clusters_tiles_hfnose_token_);
0190 const typename PatternRecognitionAlgoBaseT<TICLLayerTilesHFNose>::Inputs inputHFNose(
0191 evt, es, layerClusters, inputClusterMask, layerClustersTimes, layer_clusters_hfnose_tiles, seeding_regions);
0192
0193 myAlgoHFNose_->makeTracksters(inputHFNose, *initialResult, seedToTrackstersAssociation);
0194
0195 inferenceAlgo_->inputData(layerClusters, *initialResult);
0196 inferenceAlgo_->runInference(*initialResult);
0197 myAlgoHFNose_->filter(*result, *initialResult, inputHFNose, seedToTrackstersAssociation);
0198
0199 } else {
0200 const auto& layer_clusters_tiles = evt.get(layer_clusters_tiles_token_);
0201 const typename PatternRecognitionAlgoBaseT<TICLLayerTiles>::Inputs input(
0202 evt, es, layerClusters, inputClusterMask, layerClustersTimes, layer_clusters_tiles, seeding_regions);
0203
0204 myAlgo_->makeTracksters(input, *initialResult, seedToTrackstersAssociation);
0205
0206 inferenceAlgo_->inputData(layerClusters, *initialResult);
0207 inferenceAlgo_->runInference(*initialResult);
0208 myAlgo_->filter(*result, *initialResult, input, seedToTrackstersAssociation);
0209 }
0210 }
0211
0212 output_mask->reserve(original_layerclusters_mask.size());
0213
0214 std::copy(
0215 std::begin(original_layerclusters_mask), std::end(original_layerclusters_mask), std::back_inserter(*output_mask));
0216
0217 for (auto& trackster : *result) {
0218 trackster.setIteration(iterIndex_);
0219
0220 for (auto const v : trackster.vertices()) {
0221
0222
0223 (*output_mask)[v] = 0.;
0224 }
0225 }
0226
0227 evt.put(std::move(result));
0228 evt.put(std::move(output_mask));
0229 }