Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-09-26 05:07:09

0001 // Author: Felice Pantaleo,Marco Rovere - felice.pantaleo@cern.ch,marco.rovere@cern.ch
0002 // Date: 09/2018
0003 
0004 // user include files
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   // static methods for handling the global cache
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   // Initialize inference algorithm using the factory
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>>();  // Mask to be applied at the next iteration
0107 }
0108 
0109 void TrackstersProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0110   // hgcalMultiClusters
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", "TracksterInferenceByDNN");
0123 
0124   // CA Plugin
0125   edm::ParameterSetDescription pluginDesc;
0126   pluginDesc.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "CA", true));
0127   desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByCA", pluginDesc);
0128   //
0129   // CLUE3D Plugin
0130   edm::ParameterSetDescription pluginDescClue3D;
0131   pluginDescClue3D.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "CLUE3D", true));
0132   desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByCLUE3D", pluginDescClue3D);
0133 
0134   // FastJet Plugin
0135   edm::ParameterSetDescription pluginDescFastJet;
0136   pluginDescFastJet.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "FastJet", true));
0137   desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByFastJet", pluginDescFastJet);
0138 
0139   // PassThrough Plugin
0140   edm::ParameterSetDescription pluginDescPassThrough;
0141   pluginDescPassThrough.addNode(edm::PluginDescription<PatternRecognitionFactory>("type", "Passthrough", true));
0142   desc.add<edm::ParameterSetDescription>("pluginPatternRecognitionByPassthrough", pluginDescPassThrough);
0143 
0144   // Inference Plugins
0145   edm::ParameterSetDescription inferenceDesc;
0146   inferenceDesc.addNode(edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByDNN", true));
0147   desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByDNN", inferenceDesc);
0148 
0149   edm::ParameterSetDescription inferenceDescANN;
0150   inferenceDescANN.addNode(
0151       edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByANN", true));
0152   desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByANN", inferenceDescANN);
0153 
0154   edm::ParameterSetDescription inferenceDescCNNv4;
0155   inferenceDescCNNv4.addNode(
0156       edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByCNNv4", true));
0157   desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByCNNv4", inferenceDescCNNv4);
0158 
0159   descriptions.add("trackstersProducer", desc);
0160 }
0161 
0162 void TrackstersProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
0163   auto result = std::make_unique<std::vector<Trackster>>();
0164   auto initialResult = std::make_unique<std::vector<Trackster>>();
0165   auto output_mask = std::make_unique<std::vector<float>>();
0166 
0167   const std::vector<float>& original_layerclusters_mask = evt.get(original_layerclusters_mask_token_);
0168   const auto& layerClusters = evt.get(clusters_token_);
0169   const auto& inputClusterMask = evt.get(filtered_layerclusters_mask_token_);
0170   const auto& layerClustersTimes = evt.get(clustersTime_token_);
0171   const auto& seeding_regions = evt.get(seeding_regions_token_);
0172 
0173   std::unordered_map<int, std::vector<int>> seedToTrackstersAssociation;
0174   // if it's regional iteration and there are seeding regions
0175   if (!seeding_regions.empty() and seeding_regions[0].index != -1) {
0176     auto numberOfSeedingRegions = seeding_regions.size();
0177     for (unsigned int i = 0; i < numberOfSeedingRegions; ++i) {
0178       seedToTrackstersAssociation.emplace(seeding_regions[i].index, 0);
0179     }
0180   }
0181 
0182   if (doNose_) {
0183     const auto& layer_clusters_hfnose_tiles = evt.get(layer_clusters_tiles_hfnose_token_);
0184     const typename PatternRecognitionAlgoBaseT<TICLLayerTilesHFNose>::Inputs inputHFNose(
0185         evt, es, layerClusters, inputClusterMask, layerClustersTimes, layer_clusters_hfnose_tiles, seeding_regions);
0186 
0187     myAlgoHFNose_->makeTracksters(inputHFNose, *initialResult, seedToTrackstersAssociation);
0188     // Run inference algorithm
0189     inferenceAlgo_->inputData(layerClusters, *initialResult);
0190     inferenceAlgo_->runInference(*initialResult);
0191     myAlgoHFNose_->filter(*result, *initialResult, inputHFNose, seedToTrackstersAssociation);
0192 
0193   } else {
0194     const auto& layer_clusters_tiles = evt.get(layer_clusters_tiles_token_);
0195     const typename PatternRecognitionAlgoBaseT<TICLLayerTiles>::Inputs input(
0196         evt, es, layerClusters, inputClusterMask, layerClustersTimes, layer_clusters_tiles, seeding_regions);
0197 
0198     myAlgo_->makeTracksters(input, *initialResult, seedToTrackstersAssociation);
0199     // Run inference algorithm
0200     inferenceAlgo_->inputData(layerClusters, *initialResult);
0201     inferenceAlgo_->runInference(*initialResult);
0202     myAlgo_->filter(*result, *initialResult, input, seedToTrackstersAssociation);
0203   }
0204   // Now update the global mask and put it into the event
0205   output_mask->reserve(original_layerclusters_mask.size());
0206   // Copy over the previous state
0207   std::copy(
0208       std::begin(original_layerclusters_mask), std::end(original_layerclusters_mask), std::back_inserter(*output_mask));
0209 
0210   for (auto& trackster : *result) {
0211     trackster.setIteration(iterIndex_);
0212     // Mask the used elements, accordingly
0213     for (auto const v : trackster.vertices()) {
0214       // TODO(rovere): for the moment we mask the layer cluster completely. In
0215       // the future, properly compute the fraction of usage.
0216       (*output_mask)[v] = 0.;
0217     }
0218   }
0219 
0220   evt.put(std::move(result));
0221   evt.put(std::move(output_mask));
0222 }