Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-01-23 00:19:42

0001 #include "DataFormats/Math/interface/deltaR.h"
0002 #include <TLorentzVector.h>
0003 #include <cmath>
0004 #include <vector>
0005 
0006 #include "FWCore/Framework/interface/Frameworkfwd.h"
0007 #include "FWCore/Framework/interface/stream/EDProducer.h"
0008 #include "FWCore/Framework/interface/Event.h"
0009 #include "FWCore/Framework/interface/MakerMacros.h"
0010 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0011 
0012 #include "DataFormats/L1TParticleFlow/interface/PFTau.h"
0013 #include "DataFormats/L1TParticleFlow/interface/PFCandidate.h"
0014 #include "L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h"
0015 #include "L1Trigger/Phase2L1ParticleFlow/interface/taus/TauNNIdHW.h"
0016 
0017 #include "ap_int.h"
0018 #include "ap_fixed.h"
0019 
0020 using namespace l1t;
0021 
0022 class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<tensorflow::SessionCache>> {
0023 public:
0024   explicit L1NNTauProducer(const edm::ParameterSet&, const tensorflow::SessionCache*);
0025   ~L1NNTauProducer() override;
0026 
0027   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0028   static std::unique_ptr<tensorflow::SessionCache> initializeGlobalCache(const edm::ParameterSet&);
0029   static void globalEndJob(const tensorflow::SessionCache*){};
0030 
0031 private:
0032   // There is te software and hardware emulator for the tau, default is the Hardware.
0033   std::unique_ptr<TauNNId> fTauNNId_;
0034   std::unique_ptr<TauNNIdHW> fTauNNIdHW_;  // Default
0035 
0036   void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
0037   void process_SW(const l1t::PFCandidateCollection& parts, std::unique_ptr<l1t::PFTauCollection>& iTaus);
0038   void process_HW(const l1t::PFCandidateCollection& parts, std::unique_ptr<l1t::PFTauCollection>& iTaus);
0039   void makeTau_HW(const l1t::PFCandidate& seed,
0040                   l1t::PFCandidateCollection& parts,
0041                   std::unique_ptr<l1t::PFTauCollection>& iTaus);
0042 
0043   void addTau(const l1t::PFCandidate& iCand,
0044               const l1t::PFCandidateCollection& iParts,
0045               std::unique_ptr<PFTauCollection>& outputTaus);
0046 
0047   double fSeedPt_;
0048   double fConeSize_;
0049   double fTauSize_;
0050   int fMaxTaus_;
0051   int fNParticles_;
0052   const bool fHW;
0053   const bool fEMSeed;
0054   const bool fDebug;
0055   edm::EDGetTokenT<vector<l1t::PFCandidate>> fL1PFToken_;
0056 };
0057 
0058 static constexpr float track_trigger_eta_max = 2.5;
0059 
0060 L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const tensorflow::SessionCache* cache)
0061     : fSeedPt_(cfg.getParameter<double>("seedpt")),
0062       fConeSize_(cfg.getParameter<double>("conesize")),
0063       fTauSize_(cfg.getParameter<double>("tausize")),
0064       fMaxTaus_(cfg.getParameter<int>("maxtaus")),
0065       fNParticles_(cfg.getParameter<int>("nparticles")),
0066       fHW(cfg.getParameter<bool>("HW")),
0067       fEMSeed(cfg.getParameter<bool>("emseed")),
0068       fDebug(cfg.getParameter<bool>("debug")),
0069       fL1PFToken_(consumes<vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))) {
0070   std::string lNNFile = cfg.getParameter<std::string>("NNFileName");  //,"L1Trigger/Phase2L1Taus/data/tau_3layer.pb");
0071   if (fHW) {
0072     fTauNNIdHW_ = std::make_unique<TauNNIdHW>();
0073     fTauNNIdHW_->initialize("input_1:0", fNParticles_);
0074   } else {
0075     fTauNNId_ = std::make_unique<TauNNId>(lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0",
0076                                           cache->getSession(),
0077                                           lNNFile,
0078                                           fNParticles_);
0079   }
0080   produces<l1t::PFTauCollection>("L1PFTausNN");
0081 }
0082 
0083 std::unique_ptr<tensorflow::SessionCache> L1NNTauProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
0084   tensorflow::setLogging("3");
0085   std::string graphPath = edm::FileInPath(cfg.getParameter<std::string>("NNFileName")).fullPath();
0086   return std::make_unique<tensorflow::SessionCache>(graphPath);
0087 }
0088 
0089 void L1NNTauProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0090   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
0091   iEvent.getByToken(fL1PFToken_, l1PFCandidates);
0092   auto lTaus = std::make_unique<l1t::PFTauCollection>();
0093 
0094   if (fHW) {
0095     process_HW(*l1PFCandidates, lTaus);
0096   } else {
0097     process_SW(*l1PFCandidates, lTaus);
0098   }
0099 
0100   if (lTaus->empty()) {
0101     PFTau dummy;
0102     lTaus->push_back(dummy);
0103   }
0104   std::sort(lTaus->begin(), lTaus->end(), [](l1t::PFTau i, l1t::PFTau j) { return (i.pt() > j.pt()); });
0105   iEvent.put(std::move(lTaus), "L1PFTausNN");
0106 }
0107 void L1NNTauProducer::process_SW(const l1t::PFCandidateCollection& parts,
0108                                  std::unique_ptr<l1t::PFTauCollection>& iTaus) {
0109   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons;
0110   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons_sort_v;
0111   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons_seeds_v;
0112   for (const auto& l1PFCand : parts)
0113     if ((l1PFCand.id() == l1t::PFCandidate::ChargedHadron || l1PFCand.id() == l1t::PFCandidate::Electron) &&
0114         std::abs(l1PFCand.eta()) < track_trigger_eta_max)
0115       pfChargedHadrons_sort_v.push_back(std::make_unique<l1t::PFCandidate>(l1PFCand));
0116 
0117   if (pfChargedHadrons_sort_v.empty())
0118     return;
0119   std::sort(
0120       pfChargedHadrons_sort_v.begin(),
0121       pfChargedHadrons_sort_v.end(),
0122       [](std::unique_ptr<l1t::PFCandidate>& i, std::unique_ptr<l1t::PFCandidate>& j) { return (i->pt() > j->pt()); });
0123 
0124   pfChargedHadrons_seeds_v.push_back(std::move(pfChargedHadrons_sort_v[0]));
0125   for (unsigned int i0 = 1; i0 < pfChargedHadrons_sort_v.size(); i0++) {
0126     bool pMatch = false;
0127     for (unsigned int i1 = 0; i1 < pfChargedHadrons_seeds_v.size(); i1++) {
0128       if (reco::deltaR2(*(pfChargedHadrons_seeds_v[i1]), *(pfChargedHadrons_sort_v[i0])) < fConeSize_ * fConeSize_)
0129         pMatch = true;
0130     }
0131     if (pMatch)
0132       continue;
0133     pfChargedHadrons_seeds_v.push_back(std::move(pfChargedHadrons_sort_v[i0]));
0134     if (int(pfChargedHadrons_seeds_v.size()) > fMaxTaus_ - 1)
0135       break;
0136   }
0137   for (unsigned int i0 = 0; i0 < pfChargedHadrons_seeds_v.size(); i0++) {
0138     addTau(*(pfChargedHadrons_seeds_v[i0]), parts, iTaus);
0139   }
0140 }
0141 
0142 // create taus based on grid structure
0143 void L1NNTauProducer::addTau(const l1t::PFCandidate& iCand,
0144                              const l1t::PFCandidateCollection& iParts,
0145                              std::unique_ptr<l1t::PFTauCollection>& outputTaus) {
0146   l1t::PFCandidateCollection pfTauCands;
0147   math::PtEtaPhiMLorentzVector lTot(0, 0, 0, 0);
0148   math::PtEtaPhiMLorentzVector lCand(0, 0, 0, 0);
0149   int lId = 0;
0150   float z0 = 0;
0151   float dxy = 0;
0152   for (const auto& l1PFCand : iParts) {
0153     if (reco::deltaR2(iCand, l1PFCand) > fConeSize_ * fConeSize_)
0154       continue;
0155     math::PtEtaPhiMLorentzVector pVec(l1PFCand.pt(), l1PFCand.eta(), l1PFCand.phi(), 0);
0156     lTot += pVec;
0157     if (reco::deltaR2(iCand, l1PFCand) < fTauSize_ * fTauSize_ &&
0158         (l1PFCand.id() == l1t::PFCandidate::Electron || l1PFCand.id() == l1t::PFCandidate::ChargedHadron ||
0159          l1PFCand.id() == l1t::PFCandidate::Photon)) {
0160       lId++;
0161       lCand += pVec;
0162       if (z0 == 0 && l1PFCand.id() == l1t::PFCandidate::ChargedHadron) {
0163         z0 = l1PFCand.z0();
0164         dxy = l1PFCand.dxy();
0165       }
0166     }
0167     pfTauCands.push_back(l1PFCand);
0168   }
0169   if (lTot.Pt() < fSeedPt_)
0170     return;
0171   std::sort(
0172       pfTauCands.begin(), pfTauCands.end(), [](l1t::PFCandidate i, l1t::PFCandidate j) { return (i.pt() > j.pt()); });
0173   float NN = fTauNNId_->compute(iCand, pfTauCands);
0174   float* lNNVector = fTauNNId_->NNVectorVar();
0175   math::PtEtaPhiMLorentzVector tempP4(lCand.Pt(), lCand.Eta(), lCand.Phi(), lCand.M() * lCand.M());
0176   l1t::PFTau l1PFTau(tempP4, lNNVector, NN, 0, lId);
0177   l1PFTau.setZ0(z0);
0178   l1PFTau.setDxy(dxy);
0179   outputTaus->push_back(l1PFTau);
0180 }
0181 void L1NNTauProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0182   // L1NNTauProducer
0183   edm::ParameterSetDescription desc;
0184   desc.add<std::string>("NNFileName", "L1Trigger/Phase2L1ParticleFlow/data/tau_3layer.pb");
0185   desc.add<double>("tausize", 0.1);
0186   desc.add<int>("maxtaus", 5);
0187   desc.add<int>("nparticles", 10);
0188   desc.add<double>("conesize", 0.4);
0189   desc.add<double>("seedpt", 20);
0190   desc.add<bool>("HW", true);
0191   desc.add<bool>("emseed", true);
0192   desc.add<bool>("debug", false);
0193   desc.add<edm::InputTag>("L1PFObjects", edm::InputTag("L1PFProducer", "l1pfCandidates"));
0194   descriptions.add("L1NNTauProducer", desc);
0195 }
0196 
0197 void L1NNTauProducer::makeTau_HW(const l1t::PFCandidate& seed,
0198                                  l1t::PFCandidateCollection& parts,
0199                                  std::unique_ptr<l1t::PFTauCollection>& iTaus) {
0200   // Seed Cone Jet algorithm with ap_fixed types and hardware emulation
0201   L1TauEmu::detaphi_t rCone2 =
0202       L1TauEmu::detaphi_t(fTauSize_ * fTauSize_ * L1TauEmu::etaphi_base * L1TauEmu::etaphi_base);
0203   unsigned lId = 0;
0204 
0205   input2_t p1_tot = 0;
0206   input2_t p1x_tot = 0;
0207   input2_t p1y_tot = 0;
0208   input2_t p1z_tot = 0;
0209 
0210   input_t e1ta_1 = seed.eta();
0211   input_t p1hi_1 = seed.phi();
0212   L1TauEmu::pt_t pt = 0;
0213   L1TauEmu::z0_t z0 = 0;
0214   L1TauEmu::dxy_t dxy = 0;
0215 
0216   // Reconstruct the Tau Cone
0217   for (unsigned i0 = 0; i0 < parts.size(); i0++) {
0218     if (L1TauEmu::inCone(seed, (parts[i0]), rCone2)) {
0219       if (parts[i0].id() == l1t::PFCandidate::Electron || parts[i0].id() == l1t::PFCandidate::ChargedHadron ||
0220           parts[i0].id() == l1t::PFCandidate::Photon) {
0221         lId++;
0222         pt = pt + L1TauEmu::pt_t(parts[i0].pt());
0223 
0224         input2_t d1eta = input_t(parts[i0].eta()) - e1ta_1;
0225         input2_t d1phi = input_t(parts[i0].phi()) - p1hi_1;
0226         input2_t d1r2 = d1eta * d1eta + d1phi * d1phi;
0227         input2_t tmppt = input_t(parts[i0].pt());
0228         input2_t half = 0.5;
0229         p1z_tot = p1z_tot + tmppt * (1 - d1r2 * half);
0230         p1y_tot = p1y_tot + tmppt * d1phi;
0231         p1x_tot = p1x_tot + tmppt * d1eta;
0232         p1_tot = p1_tot + tmppt;
0233 
0234         if (z0 == 0 && parts[i0].id() == l1t::PFCandidate::ChargedHadron) {
0235           z0 = parts[i0].hwZ0();
0236           dxy = parts[i0].hwDxy();
0237         }
0238       }
0239     }
0240   }
0241 
0242   //Compute the mass
0243   input2_t tmpmass1 = (p1_tot * p1_tot - p1x_tot * p1x_tot - p1y_tot * p1y_tot - p1z_tot * p1z_tot);
0244   if (tmpmass1 < 0)
0245     tmpmass1 = 0;
0246   L1TauEmu::pt_t mass = l1ct::pt_t(tmpmass1);
0247 
0248   if (pt < fSeedPt_)
0249     return;
0250 
0251   // Tau NN Inference
0252   Tau_NN_Result NN_ouput = fTauNNIdHW_->compute(seed, parts);
0253 
0254   // Needed for making PFTau
0255   input_t* lNNVector = fTauNNIdHW_->NNVectorVar();
0256   float pNNVec[80];
0257   for (unsigned i0 = 0; i0 < 80; i0++)
0258     pNNVec[i0] = float(lNNVector[i0]);
0259 
0260   //Firmware Tau
0261   l1ct::Tau l1ctTau;
0262   l1ctTau.hwPt = l1ct::pt_t(pt * NN_ouput.nn_pt_correction);  //l1gt is <16,11> and currently <16,14>
0263   l1ctTau.hwEta = l1ct::Scales::makeGlbEta(seed.eta());       // seed.eta() and seed.phi() are in physical coordinates
0264   l1ctTau.hwPhi = l1ct::Scales::makeGlbPhi(seed.phi());
0265 
0266   l1ctTau.hwSeedPt = seed.pt();
0267   l1ctTau.hwSeedZ0 = seed.hwZ0();
0268   l1ctTau.hwCharge = seed.charge();
0269 
0270   l1ctTau.hwType = l1ct::Tau::type_t(lId);
0271   l1ctTau.hwRawId = ap_uint<10>(NN_ouput.nn_id * 1024);  //NN Output is ap_fixed<16, 6> so need to cast.
0272 
0273   //Convert to GT format and pack to encodedTau of PFTau
0274   l1gt::Tau l1gtTau = l1ctTau.toGT();
0275   l1gt::PackedTau packed_Tau = l1gtTau.pack();
0276 
0277   //Make PFTau
0278   //Save pt, eta and phi in gt scales
0279   math::PtEtaPhiMLorentzVector tempP4(l1gt::Scales::floatPt(l1gtTau.v3.pt),
0280                                       l1gt::Scales::floatEta(l1gtTau.v3.eta),
0281                                       l1gt::Scales::floatPhi(l1gtTau.v3.phi),
0282                                       float(mass));
0283 
0284   l1t::PFTau l1PFTau(tempP4, pNNVec, NN_ouput.nn_id, 0, lId);
0285   l1PFTau.setZ0(float(z0) * 0.05);    //L1TauEmu::z0_base);
0286   l1PFTau.setDxy(float(dxy) * 0.05);  //L1TauEmu::dxy_base);
0287 
0288   l1PFTau.set_encodedTau(packed_Tau);
0289 
0290   iTaus->push_back(l1PFTau);
0291 }
0292 
0293 void L1NNTauProducer::process_HW(const l1t::PFCandidateCollection& parts,
0294                                  std::unique_ptr<l1t::PFTauCollection>& iTaus) {
0295   // The fixed point algorithm emulation
0296   using namespace L1TauEmu;
0297   std::vector<l1t::PFCandidate> work;
0298   work.resize(parts.size());
0299   std::transform(parts.begin(), parts.end(), work.begin(), [](const l1t::PFCandidate& part) { return part; });
0300   std::sort(work.begin(), work.end(), [](l1t::PFCandidate i, l1t::PFCandidate j) {
0301     return (l1ct::pt_t(i.pt()) > l1ct::pt_t(j.pt()));
0302   });
0303 
0304   std::vector<l1t::PFCandidate> seeds;
0305   uint lSeed = l1t::PFCandidate::ChargedHadron;
0306   if (fEMSeed)
0307     lSeed = l1t::PFCandidate::Photon;
0308   std::copy_if(work.begin(), work.end(), std::back_inserter(seeds), [&](const l1t::PFCandidate& part) {
0309     return ((part.id() == l1t::PFCandidate::Electron || part.id() == l1t::PFCandidate::ChargedHadron ||
0310              part.id() == lSeed) &&
0311             std::abs(part.eta()) < track_trigger_eta_max);
0312   });
0313   // It would be nice to transform the inputs to the etaphi_base of the FW here, as in the line below
0314   // However the phi may wrap around if the etaphi_base > 1, so don't do it...
0315   //std::for_each(work.begin(), work.end(), [](l1t::PFCandidate& x){x.setP4(math::PtEtaPhiMLorentzVector(pt_t(x.pt()), etaphi_t(x.eta()*etaphi_base), etaphi_t(x.phi()*etaphi_base), x.mass()));});
0316   detaphi_t rCone2 = detaphi_t(fConeSize_ * fConeSize_ * etaphi_base * etaphi_base);
0317 
0318   iTaus->reserve(fMaxTaus_);
0319   while (!seeds.empty() && iTaus->size() < unsigned(fMaxTaus_)) {
0320     // Take the first (highest pt) candidate as a seed
0321     l1t::PFCandidate seed = seeds.at(0);
0322     // Get the particles within a _coneSize of the seed
0323     std::vector<l1t::PFCandidate> particlesInCone;
0324     std::copy_if(work.begin(), work.end(), std::back_inserter(particlesInCone), [&](l1t::PFCandidate& part) {
0325       return inCone(seed, part, rCone2);
0326     });
0327     makeTau_HW(seed, particlesInCone, iTaus);
0328     // remove the clustered particles
0329     work.erase(std::remove_if(
0330                    work.begin(), work.end(), [&](const l1t::PFCandidate& part) { return inCone(seed, part, rCone2); }),
0331                work.end());
0332 
0333     seeds.erase(
0334         std::remove_if(
0335             seeds.begin(), seeds.end(), [&](const l1t::PFCandidate& part) { return inCone(seed, part, rCone2); }),
0336         seeds.end());
0337   }
0338 }
0339 
0340 L1NNTauProducer::~L1NNTauProducer() {}
0341 
0342 #include "FWCore/Framework/interface/MakerMacros.h"
0343 DEFINE_FWK_MODULE(L1NNTauProducer);