Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2025-03-14 23:36:37 

0001 // Author: Felice Pantaleo, Wahid Redjeb (CERN) - felice.pantaleo@cern.ch, wahid.redjeb@cern.ch
0002 // Date: 12/2023
0003 #include <memory>  // unique_ptr
0004 #include "CommonTools/RecoAlgos/interface/MultiVectorManager.h"
0005 #include "FWCore/Framework/interface/stream/EDProducer.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0007 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0008 #include "FWCore/ParameterSet/interface/PluginDescription.h"
0009 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0010 #include "FWCore/Utilities/interface/ESGetToken.h"
0011 #include "FWCore/Framework/interface/ESHandle.h"
0012 #include "FWCore/Framework/interface/Frameworkfwd.h"
0013 #include "FWCore/Framework/interface/MakerMacros.h"
0014 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0015 #include "FWCore/Framework/interface/ConsumesCollector.h"
0016 
0017 #include "DataFormats/Common/interface/OrphanHandle.h"
0018 
0019 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0020 #include "DataFormats/HGCalReco/interface/Common.h"
0021 #include "DataFormats/HGCalReco/interface/TICLLayerTile.h"
0022 #include "DataFormats/HGCalReco/interface/Trackster.h"
0023 
0024 #include "PhysicsTools/TensorFlow/interface/TfGraphRecord.h"
0025 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0026 #include "PhysicsTools/TensorFlow/interface/TfGraphDefWrapper.h"
0027 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0028 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0029 
0030 #include "RecoHGCal/TICL/interface/TracksterLinkingAlgoBase.h"
0031 #include "RecoHGCal/TICL/plugins/TracksterLinkingPluginFactory.h"
0032 #include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h"
0033 
0034 #include "TrackingTools/GeomPropagators/interface/Propagator.h"
0035 #include "TrackingTools/Records/interface/TrackingComponentsRecord.h"
0036 
0037 #include "MagneticField/Engine/interface/MagneticField.h"
0038 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0039 
0040 #include "Geometry/HGCalCommonData/interface/HGCalDDDConstants.h"
0041 #include "Geometry/Records/interface/IdealGeometryRecord.h"
0042 #include "Geometry/CommonDetUnit/interface/GeomDet.h"
0043 
0044 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
0045 #include "RecoHGCal/TICL/interface/TracksterInferenceAlgoFactory.h"
0046 
0047 #include "TrackstersPCA.h"
0048 
0049 using namespace ticl;
0050 using cms::Ort::ONNXRuntime;
0051 
0052 class TracksterLinksProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0053 public:
0054   explicit TracksterLinksProducer(const edm::ParameterSet &ps, const ONNXRuntime *);
0055   ~TracksterLinksProducer() override {};
0056   void produce(edm::Event &, const edm::EventSetup &) override;
0057   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0058 
0059   void beginRun(edm::Run const &iEvent, edm::EventSetup const &es) override;
0060   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet &iConfig);
0061   static void globalEndJob(const ONNXRuntime *);
0062 
0063 private:
0064   void printTrackstersDebug(const std::vector<Trackster> &, const char *label) const;
0065   void dumpTrackster(const Trackster &) const;
0066 
0067   std::unique_ptr<TracksterLinkingAlgoBase> linkingAlgo_;
0068   std::string algoType_;
0069 
0070   std::vector<edm::EDGetTokenT<std::vector<Trackster>>> tracksters_tokens_;
0071   const edm::EDGetTokenT<std::vector<reco::CaloCluster>> clusters_token_;
0072   const edm::EDGetTokenT<edm::ValueMap<std::pair<float, float>>> clustersTime_token_;
0073 
0074   const bool regressionAndPid_;
0075   std::unique_ptr<TracksterInferenceAlgoBase> inferenceAlgo_;
0076 
0077   std::vector<edm::EDGetTokenT<std::vector<float>>> original_masks_tokens_;
0078 
0079   const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geometry_token_;
0080   const std::string detector_;
0081   const std::string propName_;
0082 
0083   const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> bfield_token_;
0084   const edm::ESGetToken<Propagator, TrackingComponentsRecord> propagator_token_;
0085   const HGCalDDDConstants *hgcons_;
0086   hgcal::RecHitTools rhtools_;
0087   edm::ESGetToken<HGCalDDDConstants, IdealGeometryRecord> hdc_token_;
0088 };
0089 
0090 TracksterLinksProducer::TracksterLinksProducer(const edm::ParameterSet &ps, const ONNXRuntime *onnxRuntime)
0091     : algoType_(ps.getParameter<edm::ParameterSet>("linkingPSet").getParameter<std::string>("type")),
0092       clusters_token_(consumes<std::vector<reco::CaloCluster>>(ps.getParameter<edm::InputTag>("layer_clusters"))),
0093       clustersTime_token_(
0094           consumes<edm::ValueMap<std::pair<float, float>>>(ps.getParameter<edm::InputTag>("layer_clustersTime"))),
0095       regressionAndPid_(ps.getParameter<bool>("regressionAndPid")),
0096       geometry_token_(esConsumes<CaloGeometry, CaloGeometryRecord, edm::Transition::BeginRun>()),
0097       detector_(ps.getParameter<std::string>("detector")),
0098       propName_(ps.getParameter<std::string>("propagator")),
0099       bfield_token_(esConsumes<MagneticField, IdealMagneticFieldRecord, edm::Transition::BeginRun>()),
0100       propagator_token_(
0101           esConsumes<Propagator, TrackingComponentsRecord, edm::Transition::BeginRun>(edm::ESInputTag("", propName_))) {
0102   // Loop over the edm::VInputTag and append the token to tracksters_tokens_
0103   for (auto const &tag : ps.getParameter<std::vector<edm::InputTag>>("tracksters_collections")) {
0104     tracksters_tokens_.emplace_back(consumes<std::vector<Trackster>>(tag));
0105   }
0106   //Loop over the edm::VInputTag of masks and append the token to original_masks_tokens_
0107   for (auto const &tag : ps.getParameter<std::vector<edm::InputTag>>("original_masks")) {
0108     original_masks_tokens_.emplace_back(consumes<std::vector<float>>(tag));
0109   }
0110   // Initialize inference algorithm using the factory
0111   std::string inferencePlugin = ps.getParameter<std::string>("inferenceAlgo");
0112   edm::ParameterSet inferencePSet = ps.getParameter<edm::ParameterSet>("pluginInferenceAlgo" + inferencePlugin);
0113   inferenceAlgo_ = std::unique_ptr<TracksterInferenceAlgoBase>(
0114       TracksterInferenceAlgoFactory::get()->create(inferencePlugin, inferencePSet));
0115 
0116   // New trackster collection after linking
0117   produces<std::vector<Trackster>>();
0118 
0119   // Links
0120   produces<std::vector<std::vector<unsigned int>>>();
0121   produces<std::vector<std::vector<unsigned int>>>("linkedTracksterIdToInputTracksterId");
0122   // LayerClusters Mask
0123   produces<std::vector<float>>();
0124 
0125   auto linkingPSet = ps.getParameter<edm::ParameterSet>("linkingPSet");
0126 
0127   if (algoType_ == "Skeletons") {
0128     std::string detectorName_ = (detector_ == "HFNose") ? "HGCalHFNoseSensitive" : "HGCalEESensitive";
0129     hdc_token_ = esConsumes<HGCalDDDConstants, IdealGeometryRecord, edm::Transition::BeginRun>(
0130         edm::ESInputTag("", detectorName_));
0131   }
0132 
0133   linkingAlgo_ = TracksterLinkingPluginFactory::get()->create(algoType_, linkingPSet, consumesCollector(), onnxRuntime);
0134 }
0135 
0136 std::unique_ptr<ONNXRuntime> TracksterLinksProducer::initializeGlobalCache(const edm::ParameterSet &iConfig) {
0137   auto const &pluginPset = iConfig.getParameter<edm::ParameterSet>("linkingPSet");
0138   if (pluginPset.exists("onnxModelPath"))
0139     return std::make_unique<ONNXRuntime>(pluginPset.getParameter<edm::FileInPath>("onnxModelPath").fullPath());
0140   else
0141     return std::unique_ptr<ONNXRuntime>(nullptr);
0142 }
0143 
0144 void TracksterLinksProducer::globalEndJob(const ONNXRuntime *) {}
0145 
0146 void TracksterLinksProducer::beginRun(edm::Run const &iEvent, edm::EventSetup const &es) {
0147   if (algoType_ == "Skeletons") {
0148     edm::ESHandle<HGCalDDDConstants> hdc = es.getHandle(hdc_token_);
0149     hgcons_ = hdc.product();
0150   }
0151 
0152   edm::ESHandle<CaloGeometry> geom = es.getHandle(geometry_token_);
0153   rhtools_.setGeometry(*geom);
0154 
0155   edm::ESHandle<MagneticField> bfield = es.getHandle(bfield_token_);
0156   edm::ESHandle<Propagator> propagator = es.getHandle(propagator_token_);
0157 
0158   linkingAlgo_->initialize(hgcons_, rhtools_, bfield, propagator);
0159 };
0160 
0161 void TracksterLinksProducer::dumpTrackster(const Trackster &t) const {
0162   auto e_over_h = (t.raw_em_pt() / ((t.raw_pt() - t.raw_em_pt()) != 0. ? (t.raw_pt() - t.raw_em_pt()) : 1.));
0163   LogDebug("TracksterLinksProducer")
0164       << "\nTrackster raw_pt: " << t.raw_pt() << " raw_em_pt: " << t.raw_em_pt() << " eoh: " << e_over_h
0165       << " barycenter: " << t.barycenter() << " eta,phi (baricenter): " << t.barycenter().eta() << ", "
0166       << t.barycenter().phi() << " eta,phi (eigen): " << t.eigenvectors(0).eta() << ", " << t.eigenvectors(0).phi()
0167       << " pt(eigen): " << std::sqrt(t.eigenvectors(0).Unit().perp2()) * t.raw_energy() << " seedID: " << t.seedID()
0168       << " seedIndex: " << t.seedIndex() << " size: " << t.vertices().size() << " average usage: "
0169       << (std::accumulate(std::begin(t.vertex_multiplicity()), std::end(t.vertex_multiplicity()), 0.) /
0170           (float)t.vertex_multiplicity().size())
0171       << " raw_energy: " << t.raw_energy() << " regressed energy: " << t.regressed_energy()
0172       << " probs(ga/e/mu/np/cp/nh/am/unk): ";
0173   for (auto const &p : t.id_probabilities()) {
0174     LogDebug("TracksterLinksProducer") << std::fixed << p << " ";
0175   }
0176   LogDebug("TracksterLinksProducer") << " sigmas: ";
0177   for (auto const &s : t.sigmas()) {
0178     LogDebug("TracksterLinksProducer") << s << " ";
0179   }
0180   LogDebug("TracksterLinksProducer") << std::endl;
0181 }
0182 
0183 void TracksterLinksProducer::produce(edm::Event &evt, const edm::EventSetup &es) {
0184   linkingAlgo_->setEvent(evt, es);
0185 
0186   auto resultTracksters = std::make_unique<std::vector<Trackster>>();
0187 
0188   auto linkedResultTracksters = std::make_unique<std::vector<std::vector<unsigned int>>>();
0189 
0190   const auto &layerClusters = evt.get(clusters_token_);
0191   const auto &layerClustersTimes = evt.get(clustersTime_token_);
0192 
0193   // loop over the original_masks_tokens_ and get the original masks collections and multiply them
0194   // to get the global mask
0195   std::vector<float> original_global_mask(layerClusters.size(), 1.f);
0196   for (unsigned int i = 0; i < original_masks_tokens_.size(); ++i) {
0197     const auto &tmp_mask = evt.get(original_masks_tokens_[i]);
0198     for (unsigned int j = 0; j < tmp_mask.size(); ++j) {
0199       original_global_mask[j] *= tmp_mask[j];
0200     }
0201   }
0202 
0203   auto resultMask = std::make_unique<std::vector<float>>(original_global_mask);
0204 
0205   std::vector<edm::Handle<std::vector<Trackster>>> tracksters_h(tracksters_tokens_.size());
0206   MultiVectorManager<Trackster> trackstersManager;
0207   for (unsigned int i = 0; i < tracksters_tokens_.size(); ++i) {
0208     evt.getByToken(tracksters_tokens_[i], tracksters_h[i]);
0209     //Fill MultiVectorManager
0210     trackstersManager.addVector(*tracksters_h[i]);
0211   }
0212 
0213   // Linking
0214   const typename TracksterLinkingAlgoBase::Inputs input(evt, es, layerClusters, layerClustersTimes, trackstersManager);
0215   auto linkedTracksterIdToInputTracksterId = std::make_unique<std::vector<std::vector<unsigned int>>>();
0216 
0217   // LinkTracksters will produce a vector of vector of indices of tracksters that:
0218   // 1) are linked together if more than one
0219   // 2) are isolated if only one
0220   // Result tracksters contains the final version of the trackster collection
0221   // linkedTrackstersToInputTrackstersMap contains the mapping between the linked tracksters and the input tracksters
0222   linkingAlgo_->linkTracksters(input, *resultTracksters, *linkedResultTracksters, *linkedTracksterIdToInputTracksterId);
0223 
0224   // Now we need to remove the tracksters that are not linked
0225   // We need to emplace_back in the resultTracksters only the tracksters that are linked
0226 
0227   for (auto const &resultTrackster : *resultTracksters) {
0228     for (auto const &clusterIndex : resultTrackster.vertices()) {
0229       (*resultMask)[clusterIndex] = 0.f;
0230     }
0231   }
0232 
0233   assignPCAtoTracksters(*resultTracksters,
0234                         layerClusters,
0235                         layerClustersTimes,
0236                         rhtools_.getPositionLayer(rhtools_.lastLayerEE()).z(),
0237                         rhtools_,
0238                         true);
0239 
0240   if (regressionAndPid_) {
0241     // Run inference algorithm
0242     inferenceAlgo_->inputData(layerClusters, *resultTracksters);
0243     inferenceAlgo_->runInference(
0244         *resultTracksters);  //option to use "Linking" instead of "CLU3D"/"energyAndPid" instead of "PID"
0245   }
0246 
0247   evt.put(std::move(linkedResultTracksters));
0248   evt.put(std::move(resultMask));
0249   evt.put(std::move(resultTracksters));
0250   evt.put(std::move(linkedTracksterIdToInputTracksterId), "linkedTracksterIdToInputTracksterId");
0251 }
0252 
0253 void TracksterLinksProducer::printTrackstersDebug(const std::vector<Trackster> &tracksters, const char *label) const {
0254   int counter = 0;
0255   LogDebug("TracksterLinksProducer").log([&](auto &log) {
0256     for (auto const &t : tracksters) {
0257       log << counter++ << " TracksterLinksProducer (" << label << ") obj barycenter: " << t.barycenter()
0258           << " eta,phi (baricenter): " << t.barycenter().eta() << ", " << t.barycenter().phi()
0259           << " eta,phi (eigen): " << t.eigenvectors(0).eta() << ", " << t.eigenvectors(0).phi()
0260           << " pt(eigen): " << std::sqrt(t.eigenvectors(0).Unit().perp2()) * t.raw_energy() << " seedID: " << t.seedID()
0261           << " seedIndex: " << t.seedIndex() << " size: " << t.vertices().size() << " average usage: "
0262           << (std::accumulate(std::begin(t.vertex_multiplicity()), std::end(t.vertex_multiplicity()), 0.) /
0263               (float)t.vertex_multiplicity().size())
0264           << " raw_energy: " << t.raw_energy() << " regressed energy: " << t.regressed_energy()
0265           << " probs(ga/e/mu/np/cp/nh/am/unk): ";
0266       for (auto const &p : t.id_probabilities()) {
0267         log << std::fixed << p << " ";
0268       }
0269       log << " sigmas: ";
0270       for (auto const &s : t.sigmas()) {
0271         log << s << " ";
0272       }
0273       log << "\n";
0274     }
0275   });
0276 }
0277 
0278 void TracksterLinksProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0279   edm::ParameterSetDescription desc;
0280   edm::ParameterSetDescription linkingDesc;
0281   linkingDesc.addNode(edm::PluginDescription<TracksterLinkingPluginFactory>("type", "Skeletons", true));
0282   // Inference Plugins
0283   edm::ParameterSetDescription inferenceDesc;
0284   inferenceDesc.addNode(edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByDNN", true));
0285   desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByDNN", inferenceDesc);
0286 
0287   edm::ParameterSetDescription inferenceDescPFN;
0288   inferenceDescPFN.addNode(
0289       edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByPFN", true));
0290   desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByPFN", inferenceDescPFN);
0291 
0292   edm::ParameterSetDescription inferenceDescCNNv4;
0293   inferenceDescCNNv4.addNode(
0294       edm::PluginDescription<TracksterInferenceAlgoFactory>("type", "TracksterInferenceByCNNv4", true));
0295   desc.add<edm::ParameterSetDescription>("pluginInferenceAlgoTracksterInferenceByCNNv4", inferenceDescCNNv4);
0296 
0297   desc.add<edm::ParameterSetDescription>("linkingPSet", linkingDesc);
0298   desc.add<std::vector<edm::InputTag>>("tracksters_collections", {edm::InputTag("ticlTrackstersCLUE3DHigh")});
0299   desc.add<std::vector<edm::InputTag>>("original_masks",
0300                                        {edm::InputTag("hgcalMergeLayerClusters", "InitialLayerClustersMask")});
0301   desc.add<edm::InputTag>("layer_clusters", edm::InputTag("hgcalMergeLayerClusters"));
0302   desc.add<edm::InputTag>("layer_clustersTime", edm::InputTag("hgcalMergeLayerClusters", "timeLayerCluster"));
0303   desc.add<bool>("regressionAndPid", false);
0304   desc.add<std::string>("detector", "HGCAL");
0305   desc.add<std::string>("propagator", "PropagatorWithMaterial");
0306   desc.add<std::string>("inferenceAlgo", "TracksterInferenceByPFN");
0307   descriptions.add("tracksterLinksProducer", desc);
0308 }
0309 
0310 DEFINE_FWK_MODULE(TracksterLinksProducer);
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.2.1 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!