Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:25:15

0001 #include <memory>  // unique_ptr
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0005 #include "FWCore/ParameterSet/interface/PluginDescription.h"
0006 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0007 #include "FWCore/Utilities/interface/ESGetToken.h"
0008 #include "FWCore/Framework/interface/ESHandle.h"
0009 #include "FWCore/Framework/interface/Frameworkfwd.h"
0010 #include "FWCore/Framework/interface/MakerMacros.h"
0011 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0012 #include "FWCore/Framework/interface/ConsumesCollector.h"
0013 
0014 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0015 #include "DataFormats/HGCalReco/interface/Common.h"
0016 #include "DataFormats/HGCalReco/interface/TICLLayerTile.h"
0017 #include "DataFormats/HGCalReco/interface/Trackster.h"
0018 #include "DataFormats/TrackReco/interface/Track.h"
0019 #include "DataFormats/MuonReco/interface/Muon.h"
0020 #include "DataFormats/GeometrySurface/interface/BoundDisk.h"
0021 #include "DataFormats/HGCalReco/interface/TICLCandidate.h"
0022 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0023 #include "DataFormats/Math/interface/Vector3D.h"
0024 
0025 #include "RecoHGCal/TICL/interface/GlobalCache.h"
0026 
0027 #include "PhysicsTools/TensorFlow/interface/TfGraphRecord.h"
0028 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0029 #include "PhysicsTools/TensorFlow/interface/TfGraphDefWrapper.h"
0030 
0031 #include "RecoHGCal/TICL/plugins/LinkingAlgoBase.h"
0032 #include "RecoHGCal/TICL/plugins/LinkingAlgoFactory.h"
0033 #include "RecoHGCal/TICL/plugins/LinkingAlgoByDirectionGeometric.h"
0034 
0035 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
0036 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0037 
0038 #include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h"
0039 #include "TrackingTools/GeomPropagators/interface/Propagator.h"
0040 #include "TrackingTools/Records/interface/TrackingComponentsRecord.h"
0041 
0042 #include "MagneticField/Engine/interface/MagneticField.h"
0043 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0044 
0045 #include "Geometry/HGCalCommonData/interface/HGCalDDDConstants.h"
0046 #include "Geometry/Records/interface/IdealGeometryRecord.h"
0047 #include "Geometry/CommonDetUnit/interface/GeomDet.h"
0048 
0049 #include "TrackstersPCA.h"
0050 
0051 using namespace ticl;
0052 
0053 class TrackstersMergeProducer : public edm::stream::EDProducer<> {
0054 public:
0055   explicit TrackstersMergeProducer(const edm::ParameterSet &ps);
0056   ~TrackstersMergeProducer() override{};
0057   void produce(edm::Event &, const edm::EventSetup &) override;
0058   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0059 
0060   // static methods for handling the global cache
0061   static std::unique_ptr<TrackstersCache> initializeGlobalCache(const edm::ParameterSet &);
0062   static void globalEndJob(TrackstersCache *);
0063 
0064   void beginJob();
0065   void endJob();
0066 
0067   void beginRun(edm::Run const &iEvent, edm::EventSetup const &es) override;
0068 
0069 private:
0070   typedef ticl::Trackster::IterationIndex TracksterIterIndex;
0071   typedef math::XYZVector Vector;
0072 
0073   void fillTile(TICLTracksterTiles &, const std::vector<Trackster> &, TracksterIterIndex);
0074 
0075   void energyRegressionAndID(const std::vector<reco::CaloCluster> &layerClusters,
0076                              const tensorflow::Session *,
0077                              std::vector<Trackster> &result) const;
0078   void printTrackstersDebug(const std::vector<Trackster> &, const char *label) const;
0079   void assignTimeToCandidates(std::vector<TICLCandidate> &resultCandidates) const;
0080   void dumpTrackster(const Trackster &) const;
0081 
0082   std::unique_ptr<LinkingAlgoBase> linkingAlgo_;
0083 
0084   const edm::EDGetTokenT<std::vector<Trackster>> tracksters_clue3d_token_;
0085   const edm::EDGetTokenT<std::vector<reco::CaloCluster>> clusters_token_;
0086   const edm::EDGetTokenT<edm::ValueMap<std::pair<float, float>>> clustersTime_token_;
0087   const edm::EDGetTokenT<std::vector<reco::Track>> tracks_token_;
0088   const edm::EDGetTokenT<edm::ValueMap<float>> tracks_time_token_;
0089   const edm::EDGetTokenT<edm::ValueMap<float>> tracks_time_quality_token_;
0090   const edm::EDGetTokenT<edm::ValueMap<float>> tracks_time_err_token_;
0091   const edm::EDGetTokenT<std::vector<reco::Muon>> muons_token_;
0092   const std::string tfDnnLabel_;
0093   const edm::ESGetToken<TfGraphDefWrapper, TfGraphRecord> tfDnnToken_;
0094   const tensorflow::Session *tfSession_;
0095 
0096   const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geometry_token_;
0097   const std::string detector_;
0098   const std::string propName_;
0099 
0100   const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> bfield_token_;
0101   const edm::ESGetToken<Propagator, TrackingComponentsRecord> propagator_token_;
0102   const bool optimiseAcrossTracksters_;
0103   const bool useMTDTiming_;
0104   const int eta_bin_window_;
0105   const int phi_bin_window_;
0106   const double pt_sigma_high_;
0107   const double pt_sigma_low_;
0108   const double halo_max_distance2_;
0109   const double track_min_pt_;
0110   const double track_min_eta_;
0111   const double track_max_eta_;
0112   const int track_max_missing_outerhits_;
0113   const double cosangle_align_;
0114   const double e_over_h_threshold_;
0115   const double pt_neutral_threshold_;
0116   const double resol_calo_offset_had_;
0117   const double resol_calo_scale_had_;
0118   const double resol_calo_offset_em_;
0119   const double resol_calo_scale_em_;
0120   const std::string eidInputName_;
0121   const std::string eidOutputNameEnergy_;
0122   const std::string eidOutputNameId_;
0123   const float eidMinClusterEnergy_;
0124   const int eidNLayers_;
0125   const int eidNClusters_;
0126   std::once_flag initializeGeometry_;
0127 
0128   const HGCalDDDConstants *hgcons_;
0129 
0130   std::unique_ptr<GeomDet> firstDisk_[2];
0131 
0132   tensorflow::Session *eidSession_;
0133   hgcal::RecHitTools rhtools_;
0134 
0135   static constexpr int eidNFeatures_ = 3;
0136 
0137   edm::ESGetToken<HGCalDDDConstants, IdealGeometryRecord> hdc_token_;
0138 };
0139 
0140 TrackstersMergeProducer::TrackstersMergeProducer(const edm::ParameterSet &ps)
0141     : tracksters_clue3d_token_(consumes<std::vector<Trackster>>(ps.getParameter<edm::InputTag>("trackstersclue3d"))),
0142       clusters_token_(consumes<std::vector<reco::CaloCluster>>(ps.getParameter<edm::InputTag>("layer_clusters"))),
0143       clustersTime_token_(
0144           consumes<edm::ValueMap<std::pair<float, float>>>(ps.getParameter<edm::InputTag>("layer_clustersTime"))),
0145       tracks_token_(consumes<std::vector<reco::Track>>(ps.getParameter<edm::InputTag>("tracks"))),
0146       tracks_time_token_(consumes<edm::ValueMap<float>>(ps.getParameter<edm::InputTag>("tracksTime"))),
0147       tracks_time_quality_token_(consumes<edm::ValueMap<float>>(ps.getParameter<edm::InputTag>("tracksTimeQual"))),
0148       tracks_time_err_token_(consumes<edm::ValueMap<float>>(ps.getParameter<edm::InputTag>("tracksTimeErr"))),
0149       muons_token_(consumes<std::vector<reco::Muon>>(ps.getParameter<edm::InputTag>("muons"))),
0150       tfDnnLabel_(ps.getParameter<std::string>("tfDnnLabel")),
0151       tfDnnToken_(esConsumes(edm::ESInputTag("", tfDnnLabel_))),
0152       tfSession_(nullptr),
0153       geometry_token_(esConsumes<CaloGeometry, CaloGeometryRecord, edm::Transition::BeginRun>()),
0154       detector_(ps.getParameter<std::string>("detector")),
0155       propName_(ps.getParameter<std::string>("propagator")),
0156       bfield_token_(esConsumes<MagneticField, IdealMagneticFieldRecord, edm::Transition::BeginRun>()),
0157       propagator_token_(
0158           esConsumes<Propagator, TrackingComponentsRecord, edm::Transition::BeginRun>(edm::ESInputTag("", propName_))),
0159       optimiseAcrossTracksters_(ps.getParameter<bool>("optimiseAcrossTracksters")),
0160       useMTDTiming_(ps.getParameter<bool>("useMTDTiming")),
0161       eta_bin_window_(ps.getParameter<int>("eta_bin_window")),
0162       phi_bin_window_(ps.getParameter<int>("phi_bin_window")),
0163       pt_sigma_high_(ps.getParameter<double>("pt_sigma_high")),
0164       pt_sigma_low_(ps.getParameter<double>("pt_sigma_low")),
0165       halo_max_distance2_(ps.getParameter<double>("halo_max_distance2")),
0166       track_min_pt_(ps.getParameter<double>("track_min_pt")),
0167       track_min_eta_(ps.getParameter<double>("track_min_eta")),
0168       track_max_eta_(ps.getParameter<double>("track_max_eta")),
0169       track_max_missing_outerhits_(ps.getParameter<int>("track_max_missing_outerhits")),
0170       cosangle_align_(ps.getParameter<double>("cosangle_align")),
0171       e_over_h_threshold_(ps.getParameter<double>("e_over_h_threshold")),
0172       pt_neutral_threshold_(ps.getParameter<double>("pt_neutral_threshold")),
0173       resol_calo_offset_had_(ps.getParameter<double>("resol_calo_offset_had")),
0174       resol_calo_scale_had_(ps.getParameter<double>("resol_calo_scale_had")),
0175       resol_calo_offset_em_(ps.getParameter<double>("resol_calo_offset_em")),
0176       resol_calo_scale_em_(ps.getParameter<double>("resol_calo_scale_em")),
0177       eidInputName_(ps.getParameter<std::string>("eid_input_name")),
0178       eidOutputNameEnergy_(ps.getParameter<std::string>("eid_output_name_energy")),
0179       eidOutputNameId_(ps.getParameter<std::string>("eid_output_name_id")),
0180       eidMinClusterEnergy_(ps.getParameter<double>("eid_min_cluster_energy")),
0181       eidNLayers_(ps.getParameter<int>("eid_n_layers")),
0182       eidNClusters_(ps.getParameter<int>("eid_n_clusters")),
0183       eidSession_(nullptr) {
0184   produces<std::vector<Trackster>>();
0185   produces<std::vector<TICLCandidate>>();
0186 
0187   std::string detectorName_ = (detector_ == "HFNose") ? "HGCalHFNoseSensitive" : "HGCalEESensitive";
0188   hdc_token_ =
0189       esConsumes<HGCalDDDConstants, IdealGeometryRecord, edm::Transition::BeginRun>(edm::ESInputTag("", detectorName_));
0190 
0191   auto linkingPSet = ps.getParameter<edm::ParameterSet>("linkingPSet");
0192   auto algoType = linkingPSet.getParameter<std::string>("type");
0193   linkingAlgo_ = LinkingAlgoFactory::get()->create(algoType, linkingPSet);
0194 }
0195 
0196 void TrackstersMergeProducer::beginJob() {}
0197 
0198 void TrackstersMergeProducer::endJob(){};
0199 
0200 void TrackstersMergeProducer::beginRun(edm::Run const &iEvent, edm::EventSetup const &es) {
0201   edm::ESHandle<HGCalDDDConstants> hdc = es.getHandle(hdc_token_);
0202   hgcons_ = hdc.product();
0203 
0204   edm::ESHandle<CaloGeometry> geom = es.getHandle(geometry_token_);
0205   rhtools_.setGeometry(*geom);
0206 
0207   edm::ESHandle<MagneticField> bfield = es.getHandle(bfield_token_);
0208   edm::ESHandle<Propagator> propagator = es.getHandle(propagator_token_);
0209 
0210   linkingAlgo_->initialize(hgcons_, rhtools_, bfield, propagator);
0211 };
0212 
0213 void TrackstersMergeProducer::fillTile(TICLTracksterTiles &tracksterTile,
0214                                        const std::vector<Trackster> &tracksters,
0215                                        TracksterIterIndex tracksterIteration) {
0216   int tracksterId = 0;
0217   for (auto const &t : tracksters) {
0218     tracksterTile.fill(tracksterIteration, t.barycenter().eta(), t.barycenter().phi(), tracksterId);
0219     LogDebug("TrackstersMergeProducer") << "Adding tracksterId: " << tracksterId << " into bin [eta,phi]: [ "
0220                                         << tracksterTile[tracksterIteration].etaBin(t.barycenter().eta()) << ", "
0221                                         << tracksterTile[tracksterIteration].phiBin(t.barycenter().phi())
0222                                         << "] for iteration: " << tracksterIteration << std::endl;
0223 
0224     tracksterId++;
0225   }
0226 }
0227 
0228 void TrackstersMergeProducer::dumpTrackster(const Trackster &t) const {
0229   auto e_over_h = (t.raw_em_pt() / ((t.raw_pt() - t.raw_em_pt()) != 0. ? (t.raw_pt() - t.raw_em_pt()) : 1.));
0230   LogDebug("TrackstersMergeProducer")
0231       << "\nTrackster raw_pt: " << t.raw_pt() << " raw_em_pt: " << t.raw_em_pt() << " eoh: " << e_over_h
0232       << " barycenter: " << t.barycenter() << " eta,phi (baricenter): " << t.barycenter().eta() << ", "
0233       << t.barycenter().phi() << " eta,phi (eigen): " << t.eigenvectors(0).eta() << ", " << t.eigenvectors(0).phi()
0234       << " pt(eigen): " << std::sqrt(t.eigenvectors(0).Unit().perp2()) * t.raw_energy() << " seedID: " << t.seedID()
0235       << " seedIndex: " << t.seedIndex() << " size: " << t.vertices().size() << " average usage: "
0236       << (std::accumulate(std::begin(t.vertex_multiplicity()), std::end(t.vertex_multiplicity()), 0.) /
0237           (float)t.vertex_multiplicity().size())
0238       << " raw_energy: " << t.raw_energy() << " regressed energy: " << t.regressed_energy()
0239       << " probs(ga/e/mu/np/cp/nh/am/unk): ";
0240   for (auto const &p : t.id_probabilities()) {
0241     LogDebug("TrackstersMergeProducer") << std::fixed << p << " ";
0242   }
0243   LogDebug("TrackstersMergeProducer") << " sigmas: ";
0244   for (auto const &s : t.sigmas()) {
0245     LogDebug("TrackstersMergeProducer") << s << " ";
0246   }
0247   LogDebug("TrackstersMergeProducer") << std::endl;
0248 }
0249 
0250 void TrackstersMergeProducer::produce(edm::Event &evt, const edm::EventSetup &es) {
0251   auto resultTrackstersMerged = std::make_unique<std::vector<Trackster>>();
0252   auto resultCandidates = std::make_unique<std::vector<TICLCandidate>>();
0253   auto resultFromTracks = std::make_unique<std::vector<TICLCandidate>>();
0254   tfSession_ = es.getData(tfDnnToken_).getSession();
0255 
0256   edm::Handle<std::vector<Trackster>> trackstersclue3d_h;
0257   evt.getByToken(tracksters_clue3d_token_, trackstersclue3d_h);
0258 
0259   edm::Handle<std::vector<reco::Track>> track_h;
0260   evt.getByToken(tracks_token_, track_h);
0261   const auto &tracks = *track_h;
0262 
0263   const auto &layerClusters = evt.get(clusters_token_);
0264   const auto &layerClustersTimes = evt.get(clustersTime_token_);
0265   const auto &muons = evt.get(muons_token_);
0266   edm::Handle<edm::ValueMap<float>> trackTime_h;
0267   edm::Handle<edm::ValueMap<float>> trackTimeErr_h;
0268   edm::Handle<edm::ValueMap<float>> trackTimeQual_h;
0269   if (useMTDTiming_) {
0270     evt.getByToken(tracks_time_token_, trackTime_h);
0271     evt.getByToken(tracks_time_err_token_, trackTimeErr_h);
0272     evt.getByToken(tracks_time_quality_token_, trackTimeQual_h);
0273   }
0274 
0275   // Linking
0276   linkingAlgo_->linkTracksters(track_h,
0277                                trackTime_h,
0278                                trackTimeErr_h,
0279                                trackTimeQual_h,
0280                                muons,
0281                                trackstersclue3d_h,
0282                                useMTDTiming_,
0283                                *resultCandidates,
0284                                *resultFromTracks);
0285 
0286   // Print debug info
0287   LogDebug("TrackstersMergeProducer") << "Results from the linking step : " << std::endl
0288                                       << "No. of Tracks : " << tracks.size()
0289                                       << "  No. of Tracksters : " << (*trackstersclue3d_h).size() << std::endl
0290                                       << "(neutral candidates have track id -1)" << std::endl;
0291 
0292   std::vector<TICLCandidate> &candidates = *resultCandidates;
0293   for (const auto &cand : candidates) {
0294     auto track_ptr = cand.trackPtr();
0295     auto trackster_ptrs = cand.tracksters();
0296 
0297     auto track_idx = track_ptr.get() - (edm::Ptr<reco::Track>(track_h, 0)).get();
0298     track_idx = (track_ptr.isNull()) ? -1 : track_idx;
0299 #ifdef EDM_ML_DEBUG
0300     LogDebug("TrackstersMergeProducer") << "PDG ID " << cand.pdgId() << " charge " << cand.charge() << " p " << cand.p()
0301                                         << std::endl;
0302     LogDebug("TrackstersMergeProducer") << "track id (p) : " << track_idx << " ("
0303                                         << (track_ptr.isNull() ? -1 : track_ptr->p()) << ") "
0304                                         << " trackster ids (E) : ";
0305 #endif
0306 
0307     // Merge included tracksters
0308     ticl::Trackster outTrackster;
0309     outTrackster.setTrackIdx(track_idx);
0310     auto updated_size = 0;
0311     for (const auto &ts_ptr : trackster_ptrs) {
0312 #ifdef EDM_ML_DEBUG
0313       auto ts_idx = ts_ptr.get() - (edm::Ptr<ticl::Trackster>(trackstersclue3d_h, 0)).get();
0314       LogDebug("TrackstersMergeProducer") << ts_idx << " (" << ts_ptr->raw_energy() << ") ";
0315 #endif
0316 
0317       auto &thisTrackster = *ts_ptr;
0318       updated_size += thisTrackster.vertices().size();
0319       outTrackster.vertices().reserve(updated_size);
0320       outTrackster.vertex_multiplicity().reserve(updated_size);
0321       std::copy(std::begin(thisTrackster.vertices()),
0322                 std::end(thisTrackster.vertices()),
0323                 std::back_inserter(outTrackster.vertices()));
0324       std::copy(std::begin(thisTrackster.vertex_multiplicity()),
0325                 std::end(thisTrackster.vertex_multiplicity()),
0326                 std::back_inserter(outTrackster.vertex_multiplicity()));
0327     }
0328 
0329     LogDebug("TrackstersMergeProducer") << std::endl;
0330 
0331     // Find duplicate LCs
0332     auto &orig_vtx = outTrackster.vertices();
0333     auto vtx_sorted{orig_vtx};
0334     std::sort(std::begin(vtx_sorted), std::end(vtx_sorted));
0335     for (unsigned int iLC = 1; iLC < vtx_sorted.size(); ++iLC) {
0336       if (vtx_sorted[iLC] == vtx_sorted[iLC - 1]) {
0337         // Clean up duplicate LCs
0338         const auto lcIdx = vtx_sorted[iLC];
0339         const auto firstEl = std::find(orig_vtx.begin(), orig_vtx.end(), lcIdx);
0340         const auto firstPos = std::distance(std::begin(orig_vtx), firstEl);
0341         auto iDup = std::find(std::next(firstEl), orig_vtx.end(), lcIdx);
0342         while (iDup != orig_vtx.end()) {
0343           orig_vtx.erase(iDup);
0344           outTrackster.vertex_multiplicity().erase(outTrackster.vertex_multiplicity().begin() +
0345                                                    std::distance(std::begin(orig_vtx), iDup));
0346           outTrackster.vertex_multiplicity()[firstPos] -= 1;
0347           iDup = std::find(std::next(firstEl), orig_vtx.end(), lcIdx);
0348         };
0349       }
0350     }
0351 
0352     outTrackster.zeroProbabilities();
0353     if (!outTrackster.vertices().empty()) {
0354       resultTrackstersMerged->push_back(outTrackster);
0355     }
0356   }
0357 
0358   assignPCAtoTracksters(*resultTrackstersMerged,
0359                         layerClusters,
0360                         layerClustersTimes,
0361                         rhtools_.getPositionLayer(rhtools_.lastLayerEE()).z());
0362   energyRegressionAndID(layerClusters, tfSession_, *resultTrackstersMerged);
0363 
0364   //filling the TICLCandidates information
0365   assert(resultTrackstersMerged->size() == resultCandidates->size());
0366 
0367   auto isHad = [](const Trackster &tracksterMerge) {
0368     return tracksterMerge.id_probability(Trackster::ParticleType::photon) +
0369                tracksterMerge.id_probability(Trackster::ParticleType::electron) <
0370            0.5;
0371   };
0372   for (size_t i = 0; i < resultTrackstersMerged->size(); i++) {
0373     auto const &tm = (*resultTrackstersMerged)[i];
0374     auto &cand = (*resultCandidates)[i];
0375     //common properties
0376     cand.setIdProbabilities(tm.id_probabilities());
0377     //charged candidates
0378     if (!cand.trackPtr().isNull()) {
0379       auto pdgId = isHad(tm) ? 211 : 11;
0380       auto const &tk = cand.trackPtr().get();
0381       cand.setPdgId(pdgId * tk->charge());
0382       cand.setCharge(tk->charge());
0383       cand.setRawEnergy(tm.raw_energy());
0384       auto const &regrE = tm.regressed_energy();
0385       math::XYZTLorentzVector p4(regrE * tk->momentum().unit().x(),
0386                                  regrE * tk->momentum().unit().y(),
0387                                  regrE * tk->momentum().unit().z(),
0388                                  regrE);
0389       cand.setP4(p4);
0390     } else {  // neutral candidates
0391       auto pdgId = isHad(tm) ? 130 : 22;
0392       cand.setPdgId(pdgId);
0393       cand.setCharge(0);
0394       cand.setRawEnergy(tm.raw_energy());
0395       const float &regrE = tm.regressed_energy();
0396       math::XYZTLorentzVector p4(regrE * tm.barycenter().unit().x(),
0397                                  regrE * tm.barycenter().unit().y(),
0398                                  regrE * tm.barycenter().unit().z(),
0399                                  regrE);
0400       cand.setP4(p4);
0401     }
0402   }
0403   for (auto &cand : *resultFromTracks) {  //Tracks with no linked tracksters are promoted to charged hadron candidates
0404     auto const &tk = cand.trackPtr().get();
0405     cand.setPdgId(211 * tk->charge());
0406     cand.setCharge(tk->charge());
0407     const float energy = std::sqrt(tk->p() * tk->p() + ticl::mpion2);
0408     cand.setRawEnergy(energy);
0409     math::PtEtaPhiMLorentzVector p4Polar(tk->pt(), tk->eta(), tk->phi(), ticl::mpion);
0410     cand.setP4(p4Polar);
0411   }
0412   // Compute timing
0413   resultCandidates->insert(resultCandidates->end(), resultFromTracks->begin(), resultFromTracks->end());
0414   assignTimeToCandidates(*resultCandidates);
0415 
0416   evt.put(std::move(resultTrackstersMerged));
0417   evt.put(std::move(resultCandidates));
0418 }
0419 
0420 void TrackstersMergeProducer::energyRegressionAndID(const std::vector<reco::CaloCluster> &layerClusters,
0421                                                     const tensorflow::Session *eidSession,
0422                                                     std::vector<Trackster> &tracksters) const {
0423   // Energy regression and particle identification strategy:
0424   //
0425   // 1. Set default values for regressed energy and particle id for each trackster.
0426   // 2. Store indices of tracksters whose total sum of cluster energies is above the
0427   //    eidMinClusterEnergy_ (GeV) threshold. Inference is not applied for soft tracksters.
0428   // 3. When no trackster passes the selection, return.
0429   // 4. Create input and output tensors. The batch dimension is determined by the number of
0430   //    selected tracksters.
0431   // 5. Fill input tensors with layer cluster features. Per layer, clusters are ordered descending
0432   //    by energy. Given that tensor data is contiguous in memory, we can use pointer arithmetic to
0433   //    fill values, even with batching.
0434   // 6. Zero-fill features for empty clusters in each layer.
0435   // 7. Batched inference.
0436   // 8. Assign the regressed energy and id probabilities to each trackster.
0437   //
0438   // Indices used throughout this method:
0439   // i -> batch element / trackster
0440   // j -> layer
0441   // k -> cluster
0442   // l -> feature
0443 
0444   // do nothing when no trackster passes the selection (3)
0445   int batchSize = (int)tracksters.size();
0446   if (batchSize == 0) {
0447     return;
0448   }
0449 
0450   for (auto &t : tracksters) {
0451     t.setRegressedEnergy(0.f);
0452     t.zeroProbabilities();
0453   }
0454 
0455   // create input and output tensors (4)
0456   tensorflow::TensorShape shape({batchSize, eidNLayers_, eidNClusters_, eidNFeatures_});
0457   tensorflow::Tensor input(tensorflow::DT_FLOAT, shape);
0458   tensorflow::NamedTensorList inputList = {{eidInputName_, input}};
0459   static constexpr int inputDimension = 4;
0460 
0461   std::vector<tensorflow::Tensor> outputs;
0462   std::vector<std::string> outputNames;
0463   if (!eidOutputNameEnergy_.empty()) {
0464     outputNames.push_back(eidOutputNameEnergy_);
0465   }
0466   if (!eidOutputNameId_.empty()) {
0467     outputNames.push_back(eidOutputNameId_);
0468   }
0469 
0470   // fill input tensor (5)
0471   for (int i = 0; i < batchSize; i++) {
0472     const Trackster &trackster = tracksters[i];
0473 
0474     // per layer, we only consider the first eidNClusters_ clusters in terms of
0475     // energy, so in order to avoid creating large / nested structures to do
0476     // the sorting for an unknown number of total clusters, create a sorted
0477     // list of layer cluster indices to keep track of the filled clusters
0478     std::vector<int> clusterIndices(trackster.vertices().size());
0479     for (int k = 0; k < (int)trackster.vertices().size(); k++) {
0480       clusterIndices[k] = k;
0481     }
0482     sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int &a, const int &b) {
0483       return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
0484     });
0485 
0486     // keep track of the number of seen clusters per layer
0487     std::vector<int> seenClusters(eidNLayers_);
0488 
0489     // loop through clusters by descending energy
0490     for (const int &k : clusterIndices) {
0491       // get features per layer and cluster and store the values directly in the input tensor
0492       const reco::CaloCluster &cluster = layerClusters[trackster.vertices(k)];
0493       int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
0494       if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
0495         // get the pointer to the first feature value for the current batch, layer and cluster
0496         float *features = &input.tensor<float, inputDimension>()(i, j, seenClusters[j], 0);
0497 
0498         // fill features
0499         *(features++) = float(cluster.energy() / float(trackster.vertex_multiplicity(k)));
0500         *(features++) = float(std::abs(cluster.eta()));
0501         *(features) = float(cluster.phi());
0502 
0503         // increment seen clusters
0504         seenClusters[j]++;
0505       }
0506     }
0507 
0508     // zero-fill features of empty clusters in each layer (6)
0509     for (int j = 0; j < eidNLayers_; j++) {
0510       for (int k = seenClusters[j]; k < eidNClusters_; k++) {
0511         float *features = &input.tensor<float, inputDimension>()(i, j, k, 0);
0512         for (int l = 0; l < eidNFeatures_; l++) {
0513           *(features++) = 0.f;
0514         }
0515       }
0516     }
0517   }
0518 
0519   // run the inference (7)
0520   tensorflow::run(eidSession, inputList, outputNames, &outputs);
0521 
0522   // store regressed energy per trackster (8)
0523   if (!eidOutputNameEnergy_.empty()) {
0524     // get the pointer to the energy tensor, dimension is batch x 1
0525     float *energy = outputs[0].flat<float>().data();
0526 
0527     for (int i = 0; i < batchSize; ++i) {
0528       float regressedEnergy =
0529           tracksters[i].raw_energy() > eidMinClusterEnergy_ ? energy[i] : tracksters[i].raw_energy();
0530       tracksters[i].setRegressedEnergy(regressedEnergy);
0531     }
0532   }
0533 
0534   // store id probabilities per trackster (8)
0535   if (!eidOutputNameId_.empty()) {
0536     // get the pointer to the id probability tensor, dimension is batch x id_probabilities.size()
0537     int probsIdx = !eidOutputNameEnergy_.empty();
0538     float *probs = outputs[probsIdx].flat<float>().data();
0539     int probsNumber = tracksters[0].id_probabilities().size();
0540     for (int i = 0; i < batchSize; ++i) {
0541       tracksters[i].setProbabilities(&probs[i * probsNumber]);
0542     }
0543   }
0544 }
0545 
0546 void TrackstersMergeProducer::assignTimeToCandidates(std::vector<TICLCandidate> &resultCandidates) const {
0547   for (auto &cand : resultCandidates) {
0548     if (cand.tracksters().size() > 1) {  // For single-trackster candidates the timing is already set
0549       float time = 0.f;
0550       float invTimeErr = 0.f;
0551       for (const auto &tr : cand.tracksters()) {
0552         if (tr->timeError() > 0) {
0553           auto invTimeESq = pow(tr->timeError(), -2);
0554           time += tr->time() * invTimeESq;
0555           invTimeErr += invTimeESq;
0556         }
0557       }
0558       if (invTimeErr > 0) {
0559         cand.setTime(time / invTimeErr);
0560         cand.setTimeError(sqrt(1.f / invTimeErr));
0561       }
0562     }
0563   }
0564 }
0565 
0566 void TrackstersMergeProducer::printTrackstersDebug(const std::vector<Trackster> &tracksters, const char *label) const {
0567 #ifdef EDM_ML_DEBUG
0568   int counter = 0;
0569   for (auto const &t : tracksters) {
0570     LogDebug("TrackstersMergeProducer")
0571         << counter++ << " TrackstersMergeProducer (" << label << ") obj barycenter: " << t.barycenter()
0572         << " eta,phi (baricenter): " << t.barycenter().eta() << ", " << t.barycenter().phi()
0573         << " eta,phi (eigen): " << t.eigenvectors(0).eta() << ", " << t.eigenvectors(0).phi()
0574         << " pt(eigen): " << std::sqrt(t.eigenvectors(0).Unit().perp2()) * t.raw_energy() << " seedID: " << t.seedID()
0575         << " seedIndex: " << t.seedIndex() << " size: " << t.vertices().size() << " average usage: "
0576         << (std::accumulate(std::begin(t.vertex_multiplicity()), std::end(t.vertex_multiplicity()), 0.) /
0577             (float)t.vertex_multiplicity().size())
0578         << " raw_energy: " << t.raw_energy() << " regressed energy: " << t.regressed_energy()
0579         << " probs(ga/e/mu/np/cp/nh/am/unk): ";
0580     for (auto const &p : t.id_probabilities()) {
0581       LogDebug("TrackstersMergeProducer") << std::fixed << p << " ";
0582     }
0583     LogDebug("TrackstersMergeProducer") << " sigmas: ";
0584     for (auto const &s : t.sigmas()) {
0585       LogDebug("TrackstersMergeProducer") << s << " ";
0586     }
0587     LogDebug("TrackstersMergeProducer") << std::endl;
0588   }
0589 #endif
0590 }
0591 
0592 void TrackstersMergeProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0593   edm::ParameterSetDescription desc;
0594 
0595   edm::ParameterSetDescription linkingDesc;
0596   linkingDesc.addNode(edm::PluginDescription<LinkingAlgoFactory>("type", "LinkingAlgoByDirectionGeometric", true));
0597   desc.add<edm::ParameterSetDescription>("linkingPSet", linkingDesc);
0598 
0599   desc.add<edm::InputTag>("trackstersclue3d", edm::InputTag("ticlTrackstersCLUE3DHigh"));
0600   desc.add<edm::InputTag>("layer_clusters", edm::InputTag("hgcalMergeLayerClusters"));
0601   desc.add<edm::InputTag>("layer_clustersTime", edm::InputTag("hgcalMergeLayerClusters", "timeLayerCluster"));
0602   desc.add<edm::InputTag>("tracks", edm::InputTag("generalTracks"));
0603   desc.add<edm::InputTag>("tracksTime", edm::InputTag("tofPID:t0"));
0604   desc.add<edm::InputTag>("tracksTimeQual", edm::InputTag("mtdTrackQualityMVA:mtdQualMVA"));
0605   desc.add<edm::InputTag>("tracksTimeErr", edm::InputTag("tofPID:sigmat0"));
0606   desc.add<edm::InputTag>("muons", edm::InputTag("muons1stStep"));
0607   desc.add<std::string>("detector", "HGCAL");
0608   desc.add<std::string>("propagator", "PropagatorWithMaterial");
0609   desc.add<bool>("optimiseAcrossTracksters", true);
0610   desc.add<bool>("useMTDTiming", true);
0611   desc.add<int>("eta_bin_window", 1);
0612   desc.add<int>("phi_bin_window", 1);
0613   desc.add<double>("pt_sigma_high", 2.);
0614   desc.add<double>("pt_sigma_low", 2.);
0615   desc.add<double>("halo_max_distance2", 4.);
0616   desc.add<double>("track_min_pt", 1.);
0617   desc.add<double>("track_min_eta", 1.48);
0618   desc.add<double>("track_max_eta", 3.);
0619   desc.add<int>("track_max_missing_outerhits", 5);
0620   desc.add<double>("cosangle_align", 0.9945);
0621   desc.add<double>("e_over_h_threshold", 1.);
0622   desc.add<double>("pt_neutral_threshold", 2.);
0623   desc.add<double>("resol_calo_offset_had", 1.5);
0624   desc.add<double>("resol_calo_scale_had", 0.15);
0625   desc.add<double>("resol_calo_offset_em", 1.5);
0626   desc.add<double>("resol_calo_scale_em", 0.15);
0627   desc.add<std::string>("tfDnnLabel", "tracksterSelectionTf");
0628   desc.add<std::string>("eid_input_name", "input");
0629   desc.add<std::string>("eid_output_name_energy", "output/regressed_energy");
0630   desc.add<std::string>("eid_output_name_id", "output/id_probabilities");
0631   desc.add<double>("eid_min_cluster_energy", 2.5);
0632   desc.add<int>("eid_n_layers", 50);
0633   desc.add<int>("eid_n_clusters", 10);
0634   descriptions.add("trackstersMergeProducer", desc);
0635 }
0636 
0637 DEFINE_FWK_MODULE(TrackstersMergeProducer);