Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2023-03-17 11:18:07 

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/MessageLogger/interface/MessageLogger.h"
0006 #include "FWCore/Utilities/interface/ESGetToken.h"
0007 
0008 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0009 #include "DataFormats/HGCalReco/interface/Common.h"
0010 #include "DataFormats/HGCalReco/interface/TICLLayerTile.h"
0011 #include "DataFormats/HGCalReco/interface/Trackster.h"
0012 #include "DataFormats/HGCalReco/interface/TICLSeedingRegion.h"
0013 #include "DataFormats/TrackReco/interface/Track.h"
0014 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
0015 #include "RecoHGCal/TICL/interface/GlobalCache.h"
0016 
0017 #include "PhysicsTools/TensorFlow/interface/TfGraphRecord.h"
0018 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0019 #include "PhysicsTools/TensorFlow/interface/TfGraphDefWrapper.h"
0020 
0021 #include "DataFormats/HGCalReco/interface/TICLCandidate.h"
0022 
0023 #include "TrackstersPCA.h"
0024 
0025 using namespace ticl;
0026 
0027 class TrackstersMergeProducerV3 : public edm::stream::EDProducer<> {
0028 public:
0029   explicit TrackstersMergeProducerV3(const edm::ParameterSet &ps);
0030   ~TrackstersMergeProducerV3() override{};
0031   void produce(edm::Event &, const edm::EventSetup &) override;
0032   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0033 
0034 private:
0035   typedef ticl::Trackster::IterationIndex TracksterIterIndex;
0036 
0037   void fillTile(TICLTracksterTiles &, const std::vector<Trackster> &, TracksterIterIndex);
0038 
0039   void energyRegressionAndID(const std::vector<reco::CaloCluster> &layerClusters,
0040                              const tensorflow::Session *,
0041                              std::vector<Trackster> &result) const;
0042   void printTrackstersDebug(const std::vector<Trackster> &, const char *label) const;
0043   void assignTimeToCandidates(std::vector<TICLCandidate> &resultCandidates) const;
0044   void dumpTrackster(const Trackster &) const;
0045 
0046   const edm::EDGetTokenT<std::vector<Trackster>> tracksterstrkem_token_;
0047   const edm::EDGetTokenT<std::vector<Trackster>> trackstersem_token_;
0048   const edm::EDGetTokenT<std::vector<Trackster>> tracksterstrk_token_;
0049   const edm::EDGetTokenT<std::vector<Trackster>> trackstershad_token_;
0050   const edm::EDGetTokenT<std::vector<TICLSeedingRegion>> seedingTrk_token_;
0051   const edm::EDGetTokenT<std::vector<reco::CaloCluster>> clusters_token_;
0052   const edm::EDGetTokenT<edm::ValueMap<std::pair<float, float>>> clustersTime_token_;
0053   const edm::EDGetTokenT<std::vector<reco::Track>> tracks_token_;
0054   const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geometry_token_;
0055   const std::string tfDnnLabel_;
0056   const edm::ESGetToken<TfGraphDefWrapper, TfGraphRecord> tfDnnToken_;
0057   const tensorflow::Session *tfSession_;
0058 
0059   const bool optimiseAcrossTracksters_;
0060   const int eta_bin_window_;
0061   const int phi_bin_window_;
0062   const double pt_sigma_high_;
0063   const double pt_sigma_low_;
0064   const double halo_max_distance2_;
0065   const double track_min_pt_;
0066   const double track_min_eta_;
0067   const double track_max_eta_;
0068   const int track_max_missing_outerhits_;
0069   const double cosangle_align_;
0070   const double e_over_h_threshold_;
0071   const double pt_neutral_threshold_;
0072   const double resol_calo_offset_had_;
0073   const double resol_calo_scale_had_;
0074   const double resol_calo_offset_em_;
0075   const double resol_calo_scale_em_;
0076   const std::string eidInputName_;
0077   const std::string eidOutputNameEnergy_;
0078   const std::string eidOutputNameId_;
0079   const float eidMinClusterEnergy_;
0080   const int eidNLayers_;
0081   const int eidNClusters_;
0082 
0083   tensorflow::Session *eidSession_;
0084   hgcal::RecHitTools rhtools_;
0085 
0086   static constexpr int eidNFeatures_ = 3;
0087 };
0088 
0089 TrackstersMergeProducerV3::TrackstersMergeProducerV3(const edm::ParameterSet &ps)
0090     : tracksterstrkem_token_(consumes<std::vector<Trackster>>(ps.getParameter<edm::InputTag>("tracksterstrkem"))),
0091       trackstersem_token_(consumes<std::vector<Trackster>>(ps.getParameter<edm::InputTag>("trackstersem"))),
0092       tracksterstrk_token_(consumes<std::vector<Trackster>>(ps.getParameter<edm::InputTag>("tracksterstrk"))),
0093       trackstershad_token_(consumes<std::vector<Trackster>>(ps.getParameter<edm::InputTag>("trackstershad"))),
0094       seedingTrk_token_(consumes<std::vector<TICLSeedingRegion>>(ps.getParameter<edm::InputTag>("seedingTrk"))),
0095       clusters_token_(consumes<std::vector<reco::CaloCluster>>(ps.getParameter<edm::InputTag>("layer_clusters"))),
0096       clustersTime_token_(
0097           consumes<edm::ValueMap<std::pair<float, float>>>(ps.getParameter<edm::InputTag>("layer_clustersTime"))),
0098       tracks_token_(consumes<std::vector<reco::Track>>(ps.getParameter<edm::InputTag>("tracks"))),
0099       geometry_token_(esConsumes<CaloGeometry, CaloGeometryRecord>()),
0100       tfDnnLabel_(ps.getParameter<std::string>("tfDnnLabel")),
0101       tfDnnToken_(esConsumes(edm::ESInputTag("", tfDnnLabel_))),
0102       tfSession_(nullptr),
0103       optimiseAcrossTracksters_(ps.getParameter<bool>("optimiseAcrossTracksters")),
0104       eta_bin_window_(ps.getParameter<int>("eta_bin_window")),
0105       phi_bin_window_(ps.getParameter<int>("phi_bin_window")),
0106       pt_sigma_high_(ps.getParameter<double>("pt_sigma_high")),
0107       pt_sigma_low_(ps.getParameter<double>("pt_sigma_low")),
0108       halo_max_distance2_(ps.getParameter<double>("halo_max_distance2")),
0109       track_min_pt_(ps.getParameter<double>("track_min_pt")),
0110       track_min_eta_(ps.getParameter<double>("track_min_eta")),
0111       track_max_eta_(ps.getParameter<double>("track_max_eta")),
0112       track_max_missing_outerhits_(ps.getParameter<int>("track_max_missing_outerhits")),
0113       cosangle_align_(ps.getParameter<double>("cosangle_align")),
0114       e_over_h_threshold_(ps.getParameter<double>("e_over_h_threshold")),
0115       pt_neutral_threshold_(ps.getParameter<double>("pt_neutral_threshold")),
0116       resol_calo_offset_had_(ps.getParameter<double>("resol_calo_offset_had")),
0117       resol_calo_scale_had_(ps.getParameter<double>("resol_calo_scale_had")),
0118       resol_calo_offset_em_(ps.getParameter<double>("resol_calo_offset_em")),
0119       resol_calo_scale_em_(ps.getParameter<double>("resol_calo_scale_em")),
0120       eidInputName_(ps.getParameter<std::string>("eid_input_name")),
0121       eidOutputNameEnergy_(ps.getParameter<std::string>("eid_output_name_energy")),
0122       eidOutputNameId_(ps.getParameter<std::string>("eid_output_name_id")),
0123       eidMinClusterEnergy_(ps.getParameter<double>("eid_min_cluster_energy")),
0124       eidNLayers_(ps.getParameter<int>("eid_n_layers")),
0125       eidNClusters_(ps.getParameter<int>("eid_n_clusters")),
0126       eidSession_(nullptr) {
0127   produces<std::vector<Trackster>>();
0128   produces<std::vector<TICLCandidate>>();
0129 }
0130 
0131 void TrackstersMergeProducerV3::fillTile(TICLTracksterTiles &tracksterTile,
0132                                          const std::vector<Trackster> &tracksters,
0133                                          TracksterIterIndex tracksterIteration) {
0134   int tracksterId = 0;
0135   for (auto const &t : tracksters) {
0136     tracksterTile.fill(tracksterIteration, t.barycenter().eta(), t.barycenter().phi(), tracksterId);
0137     LogDebug("TrackstersMergeProducerV3") << "Adding tracksterId: " << tracksterId << " into bin [eta,phi]: [ "
0138                                           << tracksterTile[tracksterIteration].etaBin(t.barycenter().eta()) << ", "
0139                                           << tracksterTile[tracksterIteration].phiBin(t.barycenter().phi())
0140                                           << "] for iteration: " << tracksterIteration << std::endl;
0141 
0142     tracksterId++;
0143   }
0144 }
0145 
0146 void TrackstersMergeProducerV3::dumpTrackster(const Trackster &t) const {
0147   auto e_over_h = (t.raw_em_pt() / ((t.raw_pt() - t.raw_em_pt()) != 0. ? (t.raw_pt() - t.raw_em_pt()) : 1.));
0148   LogDebug("TrackstersMergeProducerV3")
0149       << "\nTrackster raw_pt: " << t.raw_pt() << " raw_em_pt: " << t.raw_em_pt() << " eoh: " << e_over_h
0150       << " barycenter: " << t.barycenter() << " eta,phi (baricenter): " << t.barycenter().eta() << ", "
0151       << t.barycenter().phi() << " eta,phi (eigen): " << t.eigenvectors(0).eta() << ", " << t.eigenvectors(0).phi()
0152       << " pt(eigen): " << std::sqrt(t.eigenvectors(0).Unit().perp2()) * t.raw_energy() << " seedID: " << t.seedID()
0153       << " seedIndex: " << t.seedIndex() << " size: " << t.vertices().size() << " average usage: "
0154       << (std::accumulate(std::begin(t.vertex_multiplicity()), std::end(t.vertex_multiplicity()), 0.) /
0155           (float)t.vertex_multiplicity().size())
0156       << " raw_energy: " << t.raw_energy() << " regressed energy: " << t.regressed_energy()
0157       << " probs(ga/e/mu/np/cp/nh/am/unk): ";
0158   for (auto const &p : t.id_probabilities()) {
0159     LogDebug("TrackstersMergeProducerV3") << std::fixed << p << " ";
0160   }
0161   LogDebug("TrackstersMergeProducerV3") << " sigmas: ";
0162   for (auto const &s : t.sigmas()) {
0163     LogDebug("TrackstersMergeProducerV3") << s << " ";
0164   }
0165   LogDebug("TrackstersMergeProducerV3") << std::endl;
0166 }
0167 
0168 void TrackstersMergeProducerV3::produce(edm::Event &evt, const edm::EventSetup &es) {
0169   edm::ESHandle<CaloGeometry> geom = es.getHandle(geometry_token_);
0170   rhtools_.setGeometry(*geom);
0171   auto resultTrackstersMerged = std::make_unique<std::vector<Trackster>>();
0172   auto resultCandidates = std::make_unique<std::vector<TICLCandidate>>();
0173 
0174   tfSession_ = es.getData(tfDnnToken_).getSession();
0175 
0176   TICLTracksterTiles tracksterTile;
0177   std::vector<bool> usedTrackstersMerged;
0178   std::vector<int> indexInMergedCollTRKEM;
0179   std::vector<int> indexInMergedCollEM;
0180   std::vector<int> indexInMergedCollTRK;
0181   std::vector<int> indexInMergedCollHAD;
0182   std::vector<bool> usedSeeds;
0183 
0184   // associating seed to the index of the trackster in the merged collection and the iteration that found it
0185   std::map<int, std::vector<std::pair<int, TracksterIterIndex>>> seedToTracksterAssociator;
0186 
0187   edm::Handle<std::vector<reco::Track>> track_h;
0188   evt.getByToken(tracks_token_, track_h);
0189   const auto &tracks = *track_h;
0190 
0191   const auto &layerClusters = evt.get(clusters_token_);
0192   const auto &layerClustersTimes = evt.get(clustersTime_token_);
0193   const auto &trackstersEM = evt.get(trackstersem_token_);
0194   const auto &trackstersTRKEM = evt.get(tracksterstrkem_token_);
0195   const auto &trackstersTRK = evt.get(tracksterstrk_token_);
0196   const auto &trackstersHAD = evt.get(trackstershad_token_);
0197   const auto &seedingTrk = evt.get(seedingTrk_token_);
0198 
0199   usedSeeds.resize(tracks.size(), false);
0200 
0201   fillTile(tracksterTile, trackstersTRKEM, TracksterIterIndex::TRKEM);
0202   fillTile(tracksterTile, trackstersEM, TracksterIterIndex::EM);
0203   fillTile(tracksterTile, trackstersTRK, TracksterIterIndex::TRKHAD);
0204   fillTile(tracksterTile, trackstersHAD, TracksterIterIndex::HAD);
0205 
0206   auto totalNumberOfTracksters =
0207       trackstersTRKEM.size() + trackstersTRK.size() + trackstersEM.size() + trackstersHAD.size();
0208   resultTrackstersMerged->reserve(totalNumberOfTracksters);
0209   usedTrackstersMerged.resize(totalNumberOfTracksters, false);
0210   indexInMergedCollTRKEM.reserve(trackstersTRKEM.size());
0211   indexInMergedCollEM.reserve(trackstersEM.size());
0212   indexInMergedCollTRK.reserve(trackstersTRK.size());
0213   indexInMergedCollHAD.reserve(trackstersHAD.size());
0214 
0215   printTrackstersDebug(trackstersTRKEM, "tracksterTRKEM");
0216   printTrackstersDebug(trackstersEM, "tracksterEM");
0217   printTrackstersDebug(trackstersTRK, "tracksterTRK");
0218   printTrackstersDebug(trackstersHAD, "tracksterHAD");
0219 
0220   for (auto const &t : trackstersTRKEM) {
0221     indexInMergedCollTRKEM.push_back(resultTrackstersMerged->size());
0222     seedToTracksterAssociator[t.seedIndex()].emplace_back(resultTrackstersMerged->size(), TracksterIterIndex::TRKEM);
0223     resultTrackstersMerged->push_back(t);
0224   }
0225 
0226   for (auto const &t : trackstersEM) {
0227     indexInMergedCollEM.push_back(resultTrackstersMerged->size());
0228     resultTrackstersMerged->push_back(t);
0229   }
0230 
0231   for (auto const &t : trackstersTRK) {
0232     indexInMergedCollTRK.push_back(resultTrackstersMerged->size());
0233     seedToTracksterAssociator[t.seedIndex()].emplace_back(resultTrackstersMerged->size(), TracksterIterIndex::TRKHAD);
0234     resultTrackstersMerged->push_back(t);
0235   }
0236 
0237   for (auto const &t : trackstersHAD) {
0238     indexInMergedCollHAD.push_back(resultTrackstersMerged->size());
0239     resultTrackstersMerged->push_back(t);
0240   }
0241 
0242   assignPCAtoTracksters(*resultTrackstersMerged,
0243                         layerClusters,
0244                         layerClustersTimes,
0245                         rhtools_.getPositionLayer(rhtools_.lastLayerEE()).z());
0246   energyRegressionAndID(layerClusters, tfSession_, *resultTrackstersMerged);
0247 
0248   printTrackstersDebug(*resultTrackstersMerged, "TrackstersMergeProducerV3");
0249 
0250   auto trackstersMergedHandle = evt.put(std::move(resultTrackstersMerged));
0251 
0252   // TICL Candidate creation
0253   // We start from neutrals first
0254 
0255   // Photons
0256   for (unsigned i = 0; i < trackstersEM.size(); ++i) {
0257     auto mergedIdx = indexInMergedCollEM[i];
0258     usedTrackstersMerged[mergedIdx] = true;
0259     const auto &t = trackstersEM[i];  //trackster
0260     TICLCandidate tmpCandidate;
0261     tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, mergedIdx));
0262     tmpCandidate.setCharge(0);
0263     tmpCandidate.setPdgId(22);
0264     tmpCandidate.setRawEnergy(t.raw_energy());
0265     math::XYZTLorentzVector p4(t.raw_energy() * t.barycenter().unit().x(),
0266                                t.raw_energy() * t.barycenter().unit().y(),
0267                                t.raw_energy() * t.barycenter().unit().z(),
0268                                t.raw_energy());
0269     tmpCandidate.setP4(p4);
0270     resultCandidates->push_back(tmpCandidate);
0271   }
0272 
0273   // Neutral Hadrons
0274   constexpr double mpion = 0.13957;
0275   constexpr float mpion2 = mpion * mpion;
0276   for (unsigned i = 0; i < trackstersHAD.size(); ++i) {
0277     auto mergedIdx = indexInMergedCollHAD[i];
0278     usedTrackstersMerged[mergedIdx] = true;
0279     const auto &t = trackstersHAD[i];  //trackster
0280     TICLCandidate tmpCandidate;
0281     tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, mergedIdx));
0282     tmpCandidate.setCharge(0);
0283     tmpCandidate.setPdgId(130);
0284     tmpCandidate.setRawEnergy(t.raw_energy());
0285     float momentum = std::sqrt(t.raw_energy() * t.raw_energy() - mpion2);
0286     math::XYZTLorentzVector p4(momentum * t.barycenter().unit().x(),
0287                                momentum * t.barycenter().unit().y(),
0288                                momentum * t.barycenter().unit().z(),
0289                                t.raw_energy());
0290     tmpCandidate.setP4(p4);
0291     resultCandidates->push_back(tmpCandidate);
0292   }
0293 
0294   // Charged Particles
0295   for (unsigned i = 0; i < trackstersTRKEM.size(); ++i) {
0296     auto mergedIdx = indexInMergedCollTRKEM[i];
0297     if (!usedTrackstersMerged[mergedIdx]) {
0298       const auto &t = trackstersTRKEM[i];  //trackster
0299       auto trackIdx = t.seedIndex();
0300       auto const &track = tracks[trackIdx];
0301       if (!usedSeeds[trackIdx] and t.raw_energy() > 0) {
0302         usedSeeds[trackIdx] = true;
0303         usedTrackstersMerged[mergedIdx] = true;
0304 
0305         std::vector<int> trackstersTRKwithSameSeed;
0306         std::vector<int> trackstersTRKEMwithSameSeed;
0307 
0308         for (const auto &tracksterIterationPair : seedToTracksterAssociator[trackIdx]) {
0309           if (tracksterIterationPair.first != mergedIdx and !usedTrackstersMerged[tracksterIterationPair.first] and
0310               trackstersMergedHandle->at(tracksterIterationPair.first).raw_energy() > 0.) {
0311             if (tracksterIterationPair.second == TracksterIterIndex::TRKEM) {
0312               trackstersTRKEMwithSameSeed.push_back(tracksterIterationPair.first);
0313             } else if (tracksterIterationPair.second == TracksterIterIndex::TRKHAD) {
0314               trackstersTRKwithSameSeed.push_back(tracksterIterationPair.first);
0315             }
0316           }
0317         }
0318 
0319         float tracksterTotalRawPt = t.raw_pt();
0320         std::vector<int> haloTrackstersTRKIdx;
0321         bool foundCompatibleTRK = false;
0322 
0323         for (auto otherTracksterIdx : trackstersTRKwithSameSeed) {
0324           usedTrackstersMerged[otherTracksterIdx] = true;
0325           tracksterTotalRawPt += trackstersMergedHandle->at(otherTracksterIdx).raw_pt();
0326 
0327           // Check the X,Y,Z barycenter and merge if they are very close (halo)
0328           if ((t.barycenter() - trackstersMergedHandle->at(otherTracksterIdx).barycenter()).mag2() <
0329               halo_max_distance2_) {
0330             haloTrackstersTRKIdx.push_back(otherTracksterIdx);
0331 
0332           } else {
0333             foundCompatibleTRK = true;
0334           }
0335         }
0336 
0337         //check if there is 1-to-1 relationship
0338         if (trackstersTRKEMwithSameSeed.empty()) {
0339           if (foundCompatibleTRK) {
0340             TICLCandidate tmpCandidate;
0341             tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, mergedIdx));
0342             double raw_energy = t.raw_energy();
0343 
0344             tmpCandidate.setCharge(track.charge());
0345             tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, trackIdx));
0346             tmpCandidate.setPdgId(211 * track.charge());
0347             for (auto otherTracksterIdx : trackstersTRKwithSameSeed) {
0348               tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, otherTracksterIdx));
0349               raw_energy += trackstersMergedHandle->at(otherTracksterIdx).raw_energy();
0350             }
0351             tmpCandidate.setRawEnergy(raw_energy);
0352             math::XYZTLorentzVector p4(raw_energy * track.momentum().unit().x(),
0353                                        raw_energy * track.momentum().unit().y(),
0354                                        raw_energy * track.momentum().unit().z(),
0355                                        raw_energy);
0356             tmpCandidate.setP4(p4);
0357             resultCandidates->push_back(tmpCandidate);
0358 
0359           } else {
0360             TICLCandidate tmpCandidate;
0361             tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, mergedIdx));
0362             double raw_energy = t.raw_energy();
0363             tmpCandidate.setCharge(track.charge());
0364             tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, trackIdx));
0365             for (auto otherTracksterIdx : trackstersTRKwithSameSeed) {
0366               tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, otherTracksterIdx));
0367               raw_energy += trackstersMergedHandle->at(otherTracksterIdx).raw_energy();
0368             }
0369             tmpCandidate.setPdgId(11 * track.charge());
0370 
0371             tmpCandidate.setRawEnergy(raw_energy);
0372             math::XYZTLorentzVector p4(raw_energy * track.momentum().unit().x(),
0373                                        raw_energy * track.momentum().unit().y(),
0374                                        raw_energy * track.momentum().unit().z(),
0375                                        raw_energy);
0376             tmpCandidate.setP4(p4);
0377             resultCandidates->push_back(tmpCandidate);
0378           }
0379 
0380         } else {
0381           // if 1-to-many find closest trackster in momentum
0382           int closestTrackster = mergedIdx;
0383           float minPtDiff = std::abs(t.raw_pt() - track.pt());
0384           for (auto otherTracksterIdx : trackstersTRKEMwithSameSeed) {
0385             auto thisPt = tracksterTotalRawPt + trackstersMergedHandle->at(otherTracksterIdx).raw_pt() - t.raw_pt();
0386             closestTrackster = std::abs(thisPt - track.pt()) < minPtDiff ? otherTracksterIdx : closestTrackster;
0387           }
0388           usedTrackstersMerged[closestTrackster] = true;
0389 
0390           if (foundCompatibleTRK) {
0391             TICLCandidate tmpCandidate;
0392             tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, closestTrackster));
0393             double raw_energy = trackstersMergedHandle->at(closestTrackster).raw_energy();
0394 
0395             tmpCandidate.setCharge(track.charge());
0396             tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, trackIdx));
0397             tmpCandidate.setPdgId(211 * track.charge());
0398             for (auto otherTracksterIdx : trackstersTRKwithSameSeed) {
0399               tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, otherTracksterIdx));
0400               raw_energy += trackstersMergedHandle->at(otherTracksterIdx).raw_energy();
0401             }
0402             tmpCandidate.setRawEnergy(raw_energy);
0403             float momentum = std::sqrt(raw_energy * raw_energy - mpion2);
0404             math::XYZTLorentzVector p4(momentum * track.momentum().unit().x(),
0405                                        momentum * track.momentum().unit().y(),
0406                                        momentum * track.momentum().unit().z(),
0407                                        raw_energy);
0408             tmpCandidate.setP4(p4);
0409             resultCandidates->push_back(tmpCandidate);
0410 
0411           } else {
0412             TICLCandidate tmpCandidate;
0413             tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, closestTrackster));
0414             double raw_energy = trackstersMergedHandle->at(closestTrackster).raw_energy();
0415 
0416             tmpCandidate.setCharge(track.charge());
0417             tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, trackIdx));
0418             for (auto otherTracksterIdx : trackstersTRKwithSameSeed) {
0419               tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, otherTracksterIdx));
0420               raw_energy += trackstersMergedHandle->at(otherTracksterIdx).raw_energy();
0421             }
0422             tmpCandidate.setPdgId(11 * track.charge());
0423             tmpCandidate.setRawEnergy(raw_energy);
0424             math::XYZTLorentzVector p4(raw_energy * track.momentum().unit().x(),
0425                                        raw_energy * track.momentum().unit().y(),
0426                                        raw_energy * track.momentum().unit().z(),
0427                                        raw_energy);
0428             tmpCandidate.setP4(p4);
0429             resultCandidates->push_back(tmpCandidate);
0430           }
0431           // Promote all other TRKEM tracksters as photons with their energy.
0432           for (auto otherTracksterIdx : trackstersTRKEMwithSameSeed) {
0433             auto tmpIndex = (otherTracksterIdx != closestTrackster) ? otherTracksterIdx : mergedIdx;
0434             TICLCandidate photonCandidate;
0435             const auto &otherTrackster = trackstersMergedHandle->at(tmpIndex);
0436             auto gammaEnergy = otherTrackster.raw_energy();
0437             photonCandidate.setCharge(0);
0438             photonCandidate.setPdgId(22);
0439             photonCandidate.setRawEnergy(gammaEnergy);
0440             math::XYZTLorentzVector gammaP4(gammaEnergy * otherTrackster.barycenter().unit().x(),
0441                                             gammaEnergy * otherTrackster.barycenter().unit().y(),
0442                                             gammaEnergy * otherTrackster.barycenter().unit().z(),
0443                                             gammaEnergy);
0444             photonCandidate.setP4(gammaP4);
0445             photonCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, tmpIndex));
0446             resultCandidates->push_back(photonCandidate);
0447           }
0448         }
0449       }
0450     }
0451   }  //end of loop over trackstersTRKEM
0452 
0453   for (unsigned i = 0; i < trackstersTRK.size(); ++i) {
0454     auto mergedIdx = indexInMergedCollTRK[i];
0455     const auto &t = trackstersTRK[i];  //trackster
0456 
0457     if (!usedTrackstersMerged[mergedIdx] and t.raw_energy() > 0) {
0458       auto trackIdx = t.seedIndex();
0459       auto const &track = tracks[trackIdx];
0460       if (!usedSeeds[trackIdx]) {
0461         usedSeeds[trackIdx] = true;
0462         usedTrackstersMerged[mergedIdx] = true;
0463         TICLCandidate tmpCandidate;
0464         tmpCandidate.addTrackster(edm::Ptr<ticl::Trackster>(trackstersMergedHandle, mergedIdx));
0465         tmpCandidate.setCharge(track.charge());
0466         tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, trackIdx));
0467         tmpCandidate.setPdgId(211 * track.charge());
0468         tmpCandidate.setRawEnergy(t.raw_energy());
0469         float momentum = std::sqrt(t.raw_energy() * t.raw_energy() - mpion2);
0470         math::XYZTLorentzVector p4(momentum * track.momentum().unit().x(),
0471                                    momentum * track.momentum().unit().y(),
0472                                    momentum * track.momentum().unit().z(),
0473                                    t.raw_energy());
0474         tmpCandidate.setP4(p4);
0475         resultCandidates->push_back(tmpCandidate);
0476       }
0477     }
0478   }
0479   // For all seeds that have 0-energy tracksters whose track is not marked as used, create a charged hadron with the track information.
0480   for (auto const &s : seedingTrk) {
0481     if (usedSeeds[s.index] == false) {
0482       auto const &track = tracks[s.index];
0483       // emit a charged hadron
0484       TICLCandidate tmpCandidate;
0485       tmpCandidate.setCharge(track.charge());
0486       tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, s.index));
0487       tmpCandidate.setPdgId(211 * track.charge());
0488       float energy = std::sqrt(track.p() * track.p() + mpion2);
0489       tmpCandidate.setRawEnergy(energy);
0490       math::PtEtaPhiMLorentzVector p4Polar(track.pt(), track.eta(), track.phi(), mpion);
0491       tmpCandidate.setP4(p4Polar);
0492       resultCandidates->push_back(tmpCandidate);
0493       usedSeeds[s.index] = true;
0494     }
0495   }
0496 
0497   // for all general tracks (high purity, pt > 1), check if they have been used: if not, promote them as charged hadrons
0498   for (unsigned i = 0; i < tracks.size(); ++i) {
0499     auto const &track = tracks[i];
0500     if (track.pt() > track_min_pt_ and track.quality(reco::TrackBase::highPurity) and
0501         track.missingOuterHits() < track_max_missing_outerhits_ and std::abs(track.outerEta()) > track_min_eta_ and
0502         std::abs(track.outerEta()) < track_max_eta_ and usedSeeds[i] == false) {
0503       // emit a charged hadron
0504       TICLCandidate tmpCandidate;
0505       tmpCandidate.setCharge(track.charge());
0506       tmpCandidate.setTrackPtr(edm::Ptr<reco::Track>(track_h, i));
0507       tmpCandidate.setPdgId(211 * track.charge());
0508       float energy = std::sqrt(track.p() * track.p() + mpion2);
0509       tmpCandidate.setRawEnergy(energy);
0510       math::PtEtaPhiMLorentzVector p4Polar(track.pt(), track.eta(), track.phi(), mpion);
0511       tmpCandidate.setP4(p4Polar);
0512       resultCandidates->push_back(tmpCandidate);
0513       usedSeeds[i] = true;
0514     }
0515   }
0516 
0517   // Compute timing
0518   assignTimeToCandidates(*resultCandidates);
0519 
0520   evt.put(std::move(resultCandidates));
0521 }
0522 
0523 void TrackstersMergeProducerV3::energyRegressionAndID(const std::vector<reco::CaloCluster> &layerClusters,
0524                                                       const tensorflow::Session *eidSession,
0525                                                       std::vector<Trackster> &tracksters) const {
0526   // Energy regression and particle identification strategy:
0527   //
0528   // 1. Set default values for regressed energy and particle id for each trackster.
0529   // 2. Store indices of tracksters whose total sum of cluster energies is above the
0530   //    eidMinClusterEnergy_ (GeV) threshold. Inference is not applied for soft tracksters.
0531   // 3. When no trackster passes the selection, return.
0532   // 4. Create input and output tensors. The batch dimension is determined by the number of
0533   //    selected tracksters.
0534   // 5. Fill input tensors with layer cluster features. Per layer, clusters are ordered descending
0535   //    by energy. Given that tensor data is contiguous in memory, we can use pointer arithmetic to
0536   //    fill values, even with batching.
0537   // 6. Zero-fill features for empty clusters in each layer.
0538   // 7. Batched inference.
0539   // 8. Assign the regressed energy and id probabilities to each trackster.
0540   //
0541   // Indices used throughout this method:
0542   // i -> batch element / trackster
0543   // j -> layer
0544   // k -> cluster
0545   // l -> feature
0546 
0547   // set default values per trackster, determine if the cluster energy threshold is passed,
0548   // and store indices of hard tracksters
0549   std::vector<int> tracksterIndices;
0550   for (int i = 0; i < (int)tracksters.size(); i++) {
0551     // calculate the cluster energy sum (2)
0552     // note: after the loop, sumClusterEnergy might be just above the threshold
0553     // which is enough to decide whether to run inference for the trackster or
0554     // not
0555     float sumClusterEnergy = 0.;
0556     for (const unsigned int &vertex : tracksters[i].vertices()) {
0557       sumClusterEnergy += (float)layerClusters[vertex].energy();
0558       // there might be many clusters, so try to stop early
0559       if (sumClusterEnergy >= eidMinClusterEnergy_) {
0560         // set default values (1)
0561         tracksters[i].setRegressedEnergy(0.f);
0562         tracksters[i].zeroProbabilities();
0563         tracksterIndices.push_back(i);
0564         break;
0565       }
0566     }
0567   }
0568 
0569   // do nothing when no trackster passes the selection (3)
0570   int batchSize = (int)tracksterIndices.size();
0571   if (batchSize == 0) {
0572     return;
0573   }
0574 
0575   // create input and output tensors (4)
0576   tensorflow::TensorShape shape({batchSize, eidNLayers_, eidNClusters_, eidNFeatures_});
0577   tensorflow::Tensor input(tensorflow::DT_FLOAT, shape);
0578   tensorflow::NamedTensorList inputList = {{eidInputName_, input}};
0579   static constexpr int inputDimension = 4;
0580 
0581   std::vector<tensorflow::Tensor> outputs;
0582   std::vector<std::string> outputNames;
0583   if (!eidOutputNameEnergy_.empty()) {
0584     outputNames.push_back(eidOutputNameEnergy_);
0585   }
0586   if (!eidOutputNameId_.empty()) {
0587     outputNames.push_back(eidOutputNameId_);
0588   }
0589 
0590   // fill input tensor (5)
0591   for (int i = 0; i < batchSize; i++) {
0592     const Trackster &trackster = tracksters[tracksterIndices[i]];
0593 
0594     // per layer, we only consider the first eidNClusters_ clusters in terms of
0595     // energy, so in order to avoid creating large / nested structures to do
0596     // the sorting for an unknown number of total clusters, create a sorted
0597     // list of layer cluster indices to keep track of the filled clusters
0598     std::vector<int> clusterIndices(trackster.vertices().size());
0599     for (int k = 0; k < (int)trackster.vertices().size(); k++) {
0600       clusterIndices[k] = k;
0601     }
0602     sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int &a, const int &b) {
0603       return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
0604     });
0605 
0606     // keep track of the number of seen clusters per layer
0607     std::vector<int> seenClusters(eidNLayers_);
0608 
0609     // loop through clusters by descending energy
0610     for (const int &k : clusterIndices) {
0611       // get features per layer and cluster and store the values directly in the input tensor
0612       const reco::CaloCluster &cluster = layerClusters[trackster.vertices(k)];
0613       int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
0614       if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
0615         // get the pointer to the first feature value for the current batch, layer and cluster
0616         float *features = &input.tensor<float, inputDimension>()(i, j, seenClusters[j], 0);
0617 
0618         // fill features
0619         *(features++) = float(cluster.energy() / float(trackster.vertex_multiplicity(k)));
0620         *(features++) = float(std::abs(cluster.eta()));
0621         *(features) = float(cluster.phi());
0622 
0623         // increment seen clusters
0624         seenClusters[j]++;
0625       }
0626     }
0627 
0628     // zero-fill features of empty clusters in each layer (6)
0629     for (int j = 0; j < eidNLayers_; j++) {
0630       for (int k = seenClusters[j]; k < eidNClusters_; k++) {
0631         float *features = &input.tensor<float, inputDimension>()(i, j, k, 0);
0632         for (int l = 0; l < eidNFeatures_; l++) {
0633           *(features++) = 0.f;
0634         }
0635       }
0636     }
0637   }
0638 
0639   // run the inference (7)
0640   tensorflow::run(eidSession, inputList, outputNames, &outputs);
0641 
0642   // store regressed energy per trackster (8)
0643   if (!eidOutputNameEnergy_.empty()) {
0644     // get the pointer to the energy tensor, dimension is batch x 1
0645     float *energy = outputs[0].flat<float>().data();
0646 
0647     for (const int &i : tracksterIndices) {
0648       tracksters[i].setRegressedEnergy(*(energy++));
0649     }
0650   }
0651 
0652   // store id probabilities per trackster (8)
0653   if (!eidOutputNameId_.empty()) {
0654     // get the pointer to the id probability tensor, dimension is batch x id_probabilities.size()
0655     int probsIdx = !eidOutputNameEnergy_.empty();
0656     float *probs = outputs[probsIdx].flat<float>().data();
0657 
0658     for (const int &i : tracksterIndices) {
0659       tracksters[i].setProbabilities(probs);
0660       probs += tracksters[i].id_probabilities().size();
0661     }
0662   }
0663 }
0664 
0665 void TrackstersMergeProducerV3::assignTimeToCandidates(std::vector<TICLCandidate> &resultCandidates) const {
0666   for (auto &cand : resultCandidates) {
0667     if (cand.tracksters().size() > 1) {  // For single-trackster candidates the timing is already set
0668       float time = 0.f;
0669       float timeErr = 0.f;
0670       for (const auto &tr : cand.tracksters()) {
0671         if (tr->timeError() > 0) {
0672           auto invTimeESq = pow(tr->timeError(), -2);
0673           time += tr->time() * invTimeESq;
0674           timeErr += invTimeESq;
0675         }
0676       }
0677       if (timeErr > 0) {
0678         timeErr = 1. / timeErr;
0679 
0680         cand.setTime(time * timeErr);
0681         cand.setTimeError(sqrt(timeErr));
0682       }
0683     }
0684   }
0685 }
0686 
0687 void TrackstersMergeProducerV3::printTrackstersDebug(const std::vector<Trackster> &tracksters,
0688                                                      const char *label) const {
0689 #ifdef EDM_ML_DEBUG
0690   int counter = 0;
0691   for (auto const &t : tracksters) {
0692     LogDebug("TrackstersMergeProducerV3")
0693         << counter++ << " TrackstersMergeProducerV3 (" << label << ") obj barycenter: " << t.barycenter()
0694         << " eta,phi (baricenter): " << t.barycenter().eta() << ", " << t.barycenter().phi()
0695         << " eta,phi (eigen): " << t.eigenvectors(0).eta() << ", " << t.eigenvectors(0).phi()
0696         << " pt(eigen): " << std::sqrt(t.eigenvectors(0).Unit().perp2()) * t.raw_energy() << " seedID: " << t.seedID()
0697         << " seedIndex: " << t.seedIndex() << " size: " << t.vertices().size() << " average usage: "
0698         << (std::accumulate(std::begin(t.vertex_multiplicity()), std::end(t.vertex_multiplicity()), 0.) /
0699             (float)t.vertex_multiplicity().size())
0700         << " raw_energy: " << t.raw_energy() << " regressed energy: " << t.regressed_energy()
0701         << " probs(ga/e/mu/np/cp/nh/am/unk): ";
0702     for (auto const &p : t.id_probabilities()) {
0703       LogDebug("TrackstersMergeProducerV3") << std::fixed << p << " ";
0704     }
0705     LogDebug("TrackstersMergeProducerV3") << " sigmas: ";
0706     for (auto const &s : t.sigmas()) {
0707       LogDebug("TrackstersMergeProducerV3") << s << " ";
0708     }
0709     LogDebug("TrackstersMergeProducerV3") << std::endl;
0710   }
0711 #endif
0712 }
0713 
0714 void TrackstersMergeProducerV3::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0715   edm::ParameterSetDescription desc;
0716   desc.add<edm::InputTag>("tracksterstrkem", edm::InputTag("ticlTrackstersTrkEM"));
0717   desc.add<edm::InputTag>("trackstersem", edm::InputTag("ticlTrackstersEM"));
0718   desc.add<edm::InputTag>("tracksterstrk", edm::InputTag("ticlTrackstersTrk"));
0719   desc.add<edm::InputTag>("trackstershad", edm::InputTag("ticlTrackstersHAD"));
0720   desc.add<edm::InputTag>("seedingTrk", edm::InputTag("ticlSeedingTrk"));
0721   desc.add<edm::InputTag>("layer_clusters", edm::InputTag("hgcalLayerClusters"));
0722   desc.add<edm::InputTag>("layer_clustersTime", edm::InputTag("hgcalLayerClusters", "timeLayerCluster"));
0723   desc.add<edm::InputTag>("tracks", edm::InputTag("generalTracks"));
0724   desc.add<bool>("optimiseAcrossTracksters", true);
0725   desc.add<int>("eta_bin_window", 1);
0726   desc.add<int>("phi_bin_window", 1);
0727   desc.add<double>("pt_sigma_high", 2.);
0728   desc.add<double>("pt_sigma_low", 2.);
0729   desc.add<double>("halo_max_distance2", 4.);
0730   desc.add<double>("track_min_pt", 1.);
0731   desc.add<double>("track_min_eta", 1.48);
0732   desc.add<double>("track_max_eta", 3.);
0733   desc.add<int>("track_max_missing_outerhits", 5);
0734   desc.add<double>("cosangle_align", 0.9945);
0735   desc.add<double>("e_over_h_threshold", 1.);
0736   desc.add<double>("pt_neutral_threshold", 2.);
0737   desc.add<double>("resol_calo_offset_had", 1.5);
0738   desc.add<double>("resol_calo_scale_had", 0.15);
0739   desc.add<double>("resol_calo_offset_em", 1.5);
0740   desc.add<double>("resol_calo_scale_em", 0.15);
0741   desc.add<std::string>("tfDnnLabel", "tracksterSelectionTf");
0742   desc.add<std::string>("eid_input_name", "input");
0743   desc.add<std::string>("eid_output_name_energy", "output/regressed_energy");
0744   desc.add<std::string>("eid_output_name_id", "output/id_probabilities");
0745   desc.add<double>("eid_min_cluster_energy", 1.);
0746   desc.add<int>("eid_n_layers", 50);
0747   desc.add<int>("eid_n_clusters", 10);
0748   descriptions.add("trackstersMergeProducerV3", desc);
0749 }
0750 
0751 #include "FWCore/Framework/interface/MakerMacros.h"
0752 DEFINE_FWK_MODULE(TrackstersMergeProducerV3);
[ 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!