Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-03-17 11:13:13

0001 #include "DataFormats/Math/interface/deltaR.h"
0002 #include <cmath>
0003 #include <vector>
0004 
0005 #include "FWCore/Framework/interface/Frameworkfwd.h"
0006 #include "FWCore/Framework/interface/stream/EDProducer.h"
0007 #include "FWCore/Framework/interface/Event.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0010 
0011 #include "DataFormats/L1TParticleFlow/interface/PFTau.h"
0012 #include "DataFormats/L1TParticleFlow/interface/PFCandidate.h"
0013 #include "L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h"
0014 
0015 using namespace l1t;
0016 
0017 class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<tensorflow::SessionCache>> {
0018 public:
0019   explicit L1NNTauProducer(const edm::ParameterSet&, const tensorflow::SessionCache*);
0020   ~L1NNTauProducer() override;
0021 
0022   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0023   static std::unique_ptr<tensorflow::SessionCache> initializeGlobalCache(const edm::ParameterSet&);
0024   static void globalEndJob(const tensorflow::SessionCache*){};
0025 
0026 private:
0027   std::unique_ptr<TauNNId> fTauNNId_;
0028   void addTau(const l1t::PFCandidate& iCand,
0029               const l1t::PFCandidateCollection& iParts,
0030               std::unique_ptr<PFTauCollection>& outputTaus);
0031   void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
0032 
0033   double fSeedPt_;
0034   double fConeSize_;
0035   double fTauSize_;
0036   int fMaxTaus_;
0037   int fNParticles_;
0038   edm::EDGetTokenT<vector<l1t::PFCandidate>> fL1PFToken_;
0039 };
0040 
0041 static constexpr float track_trigger_eta_max = 2.5;
0042 
0043 L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const tensorflow::SessionCache* cache)
0044     : fSeedPt_(cfg.getParameter<double>("seedpt")),
0045       fConeSize_(cfg.getParameter<double>("conesize")),
0046       fTauSize_(cfg.getParameter<double>("tausize")),
0047       fMaxTaus_(cfg.getParameter<int>("maxtaus")),
0048       fNParticles_(cfg.getParameter<int>("nparticles")),
0049       fL1PFToken_(consumes<vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))) {
0050   std::string lNNFile = cfg.getParameter<std::string>("NNFileName");  //,"L1Trigger/Phase2L1Taus/data/tau_3layer.pb");
0051   fTauNNId_ = std::make_unique<TauNNId>(lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0",
0052                                         cache->getSession(),
0053                                         lNNFile,
0054                                         fNParticles_);
0055   produces<l1t::PFTauCollection>("L1PFTausNN");
0056 }
0057 
0058 std::unique_ptr<tensorflow::SessionCache> L1NNTauProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
0059   tensorflow::setLogging("3");
0060   std::string graphPath = edm::FileInPath(cfg.getParameter<std::string>("NNFileName")).fullPath();
0061   return std::make_unique<tensorflow::SessionCache>(graphPath);
0062 }
0063 
0064 void L1NNTauProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0065   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
0066   iEvent.getByToken(fL1PFToken_, l1PFCandidates);
0067 
0068   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons_sort_v;
0069   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons_seeds_v;
0070   for (const auto& l1PFCand : *l1PFCandidates)
0071     if ((l1PFCand.id() == l1t::PFCandidate::ChargedHadron || l1PFCand.id() == l1t::PFCandidate::Electron) &&
0072         std::abs(l1PFCand.eta()) < track_trigger_eta_max)
0073       pfChargedHadrons_sort_v.push_back(std::make_unique<l1t::PFCandidate>(l1PFCand));
0074 
0075   std::sort(
0076       pfChargedHadrons_sort_v.begin(),
0077       pfChargedHadrons_sort_v.end(),
0078       [](std::unique_ptr<l1t::PFCandidate>& i, std::unique_ptr<l1t::PFCandidate>& j) { return (i->pt() > j->pt()); });
0079 
0080   auto lTaus = std::make_unique<l1t::PFTauCollection>();
0081   if (pfChargedHadrons_sort_v.empty()) {
0082     if (lTaus->empty()) {
0083       PFTau dummy;
0084       lTaus->push_back(dummy);
0085     }
0086     iEvent.put(std::move(lTaus), "L1PFTausNN");
0087     return;
0088   }
0089   pfChargedHadrons_seeds_v.push_back(std::move(pfChargedHadrons_sort_v[0]));
0090   for (unsigned int i0 = 1; i0 < pfChargedHadrons_sort_v.size(); i0++) {
0091     bool pMatch = false;
0092     for (unsigned int i1 = 0; i1 < pfChargedHadrons_seeds_v.size(); i1++) {
0093       if (reco::deltaR2(*(pfChargedHadrons_seeds_v[i1]), *(pfChargedHadrons_sort_v[i0])) < fConeSize_ * fConeSize_)
0094         pMatch = true;
0095     }
0096     if (pMatch)
0097       continue;
0098     pfChargedHadrons_seeds_v.push_back(std::move(pfChargedHadrons_sort_v[i0]));
0099     if (int(pfChargedHadrons_seeds_v.size()) > fMaxTaus_ - 1)
0100       break;
0101   }
0102   for (unsigned int i0 = 0; i0 < pfChargedHadrons_seeds_v.size(); i0++) {
0103     addTau(*(pfChargedHadrons_seeds_v[i0]), (*l1PFCandidates), lTaus);
0104   }
0105   if (lTaus->empty()) {
0106     PFTau dummy;
0107     lTaus->push_back(dummy);
0108   }
0109   std::sort(lTaus->begin(), lTaus->end(), [](l1t::PFTau i, l1t::PFTau j) { return (i.pt() > j.pt()); });
0110   iEvent.put(std::move(lTaus), "L1PFTausNN");
0111 }
0112 
0113 // create taus based on grid structure
0114 void L1NNTauProducer::addTau(const l1t::PFCandidate& iCand,
0115                              const l1t::PFCandidateCollection& iParts,
0116                              std::unique_ptr<l1t::PFTauCollection>& outputTaus) {
0117   l1t::PFCandidateCollection pfTauCands;
0118   math::PtEtaPhiMLorentzVector lTot(0, 0, 0, 0);
0119   math::PtEtaPhiMLorentzVector lCand(0, 0, 0, 0);
0120   int lId = 0;
0121   for (const auto& l1PFCand : iParts) {
0122     if (reco::deltaR2(iCand, l1PFCand) > fConeSize_ * fConeSize_)
0123       continue;
0124     math::PtEtaPhiMLorentzVector pVec(l1PFCand.pt(), l1PFCand.eta(), l1PFCand.phi(), 0);
0125     lTot += pVec;
0126     if (reco::deltaR2(iCand, l1PFCand) < fTauSize_ * fTauSize_ &&
0127         (l1PFCand.id() == l1t::PFCandidate::Electron || l1PFCand.id() == l1t::PFCandidate::ChargedHadron ||
0128          l1PFCand.id() == l1t::PFCandidate::Photon)) {
0129       lId++;
0130       lCand += pVec;
0131     }
0132     pfTauCands.push_back(l1PFCand);
0133   }
0134   if (lTot.Pt() < fSeedPt_)
0135     return;
0136   std::sort(
0137       pfTauCands.begin(), pfTauCands.end(), [](l1t::PFCandidate i, l1t::PFCandidate j) { return (i.pt() > j.pt()); });
0138   float NN = fTauNNId_->compute(iCand, pfTauCands);
0139   math::PtEtaPhiMLorentzVector tempP4(lCand.Pt(), lCand.Eta(), lCand.Phi(), lCand.M());
0140   l1t::PFTau l1PFTau(tempP4, NN, 0, lId);
0141   outputTaus->push_back(l1PFTau);
0142 }
0143 void L1NNTauProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0144   // L1NNTauProducer
0145   edm::ParameterSetDescription desc;
0146   desc.add<std::string>("NNFileName", "L1Trigger/Phase2L1ParticleFlow/data/tau_3layer.pb");
0147   desc.add<double>("tausize", 0.1);
0148   desc.add<int>("maxtaus", 5);
0149   desc.add<int>("nparticles", 10);
0150   desc.add<double>("conesize", 0.4);
0151   desc.add<double>("seedpt", 20);
0152   desc.add<edm::InputTag>("L1PFObjects", edm::InputTag("L1PFProducer", "l1pfCandidates"));
0153   descriptions.add("L1NNTauProducer", desc);
0154 }
0155 L1NNTauProducer::~L1NNTauProducer() {}
0156 
0157 #include "FWCore/Framework/interface/MakerMacros.h"
0158 DEFINE_FWK_MODULE(L1NNTauProducer);