Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-09-24 22:51:30

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   std::string graphPath = edm::FileInPath(cfg.getParameter<std::string>("NNFileName")).fullPath();
0085   return std::make_unique<tensorflow::SessionCache>(graphPath);
0086 }
0087 
0088 void L1NNTauProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0089   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
0090   iEvent.getByToken(fL1PFToken_, l1PFCandidates);
0091   auto lTaus = std::make_unique<l1t::PFTauCollection>();
0092 
0093   if (fHW) {
0094     process_HW(*l1PFCandidates, lTaus);
0095   } else {
0096     process_SW(*l1PFCandidates, lTaus);
0097   }
0098 
0099   std::sort(lTaus->begin(), lTaus->end(), [](l1t::PFTau i, l1t::PFTau j) { return (i.pt() > j.pt()); });
0100   iEvent.put(std::move(lTaus), "L1PFTausNN");
0101 }
0102 void L1NNTauProducer::process_SW(const l1t::PFCandidateCollection& parts,
0103                                  std::unique_ptr<l1t::PFTauCollection>& iTaus) {
0104   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons;
0105   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons_sort_v;
0106   std::vector<unique_ptr<l1t::PFCandidate>> pfChargedHadrons_seeds_v;
0107   for (const auto& l1PFCand : parts)
0108     if ((l1PFCand.id() == l1t::PFCandidate::ChargedHadron || l1PFCand.id() == l1t::PFCandidate::Electron) &&
0109         std::abs(l1PFCand.eta()) < track_trigger_eta_max)
0110       pfChargedHadrons_sort_v.push_back(std::make_unique<l1t::PFCandidate>(l1PFCand));
0111 
0112   if (pfChargedHadrons_sort_v.empty())
0113     return;
0114   std::sort(
0115       pfChargedHadrons_sort_v.begin(),
0116       pfChargedHadrons_sort_v.end(),
0117       [](std::unique_ptr<l1t::PFCandidate>& i, std::unique_ptr<l1t::PFCandidate>& j) { return (i->pt() > j->pt()); });
0118 
0119   pfChargedHadrons_seeds_v.push_back(std::move(pfChargedHadrons_sort_v[0]));
0120   for (unsigned int i0 = 1; i0 < pfChargedHadrons_sort_v.size(); i0++) {
0121     bool pMatch = false;
0122     for (unsigned int i1 = 0; i1 < pfChargedHadrons_seeds_v.size(); i1++) {
0123       if (reco::deltaR2(*(pfChargedHadrons_seeds_v[i1]), *(pfChargedHadrons_sort_v[i0])) < fConeSize_ * fConeSize_)
0124         pMatch = true;
0125     }
0126     if (pMatch)
0127       continue;
0128     pfChargedHadrons_seeds_v.push_back(std::move(pfChargedHadrons_sort_v[i0]));
0129     if (int(pfChargedHadrons_seeds_v.size()) > fMaxTaus_ - 1)
0130       break;
0131   }
0132   for (unsigned int i0 = 0; i0 < pfChargedHadrons_seeds_v.size(); i0++) {
0133     addTau(*(pfChargedHadrons_seeds_v[i0]), parts, iTaus);
0134   }
0135 }
0136 
0137 // create taus based on grid structure
0138 void L1NNTauProducer::addTau(const l1t::PFCandidate& iCand,
0139                              const l1t::PFCandidateCollection& iParts,
0140                              std::unique_ptr<l1t::PFTauCollection>& outputTaus) {
0141   l1t::PFCandidateCollection pfTauCands;
0142   math::PtEtaPhiMLorentzVector lTot(0, 0, 0, 0);
0143   math::PtEtaPhiMLorentzVector lCand(0, 0, 0, 0);
0144   int lId = 0;
0145   float z0 = 0;
0146   float dxy = 0;
0147   for (const auto& l1PFCand : iParts) {
0148     if (reco::deltaR2(iCand, l1PFCand) > fConeSize_ * fConeSize_)
0149       continue;
0150     math::PtEtaPhiMLorentzVector pVec(l1PFCand.pt(), l1PFCand.eta(), l1PFCand.phi(), 0);
0151     lTot += pVec;
0152     if (reco::deltaR2(iCand, l1PFCand) < fTauSize_ * fTauSize_ &&
0153         (l1PFCand.id() == l1t::PFCandidate::Electron || l1PFCand.id() == l1t::PFCandidate::ChargedHadron ||
0154          l1PFCand.id() == l1t::PFCandidate::Photon)) {
0155       lId++;
0156       lCand += pVec;
0157       if (z0 == 0 && l1PFCand.id() == l1t::PFCandidate::ChargedHadron) {
0158         z0 = l1PFCand.z0();
0159         dxy = l1PFCand.dxy();
0160       }
0161     }
0162     pfTauCands.push_back(l1PFCand);
0163   }
0164   if (lTot.Pt() < fSeedPt_)
0165     return;
0166   std::sort(
0167       pfTauCands.begin(), pfTauCands.end(), [](l1t::PFCandidate i, l1t::PFCandidate j) { return (i.pt() > j.pt()); });
0168   float NN = fTauNNId_->compute(iCand, pfTauCands);
0169   float* lNNVector = fTauNNId_->NNVectorVar();
0170   math::PtEtaPhiMLorentzVector tempP4(lCand.Pt(), lCand.Eta(), lCand.Phi(), lCand.M() * lCand.M());
0171   l1t::PFTau l1PFTau(tempP4, lNNVector, NN, 0, lId);
0172   l1PFTau.setZ0(z0);
0173   l1PFTau.setDxy(dxy);
0174   outputTaus->push_back(l1PFTau);
0175 }
0176 void L1NNTauProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0177   // L1NNTauProducer
0178   edm::ParameterSetDescription desc;
0179   desc.add<std::string>("NNFileName", "L1Trigger/Phase2L1ParticleFlow/data/tau_3layer.pb");
0180   desc.add<double>("tausize", 0.1);
0181   desc.add<int>("maxtaus", 5);
0182   desc.add<int>("nparticles", 10);
0183   desc.add<double>("conesize", 0.4);
0184   desc.add<double>("seedpt", 20);
0185   desc.add<bool>("HW", true);
0186   desc.add<bool>("emseed", true);
0187   desc.add<bool>("debug", false);
0188   desc.add<edm::InputTag>("L1PFObjects", edm::InputTag("L1PFProducer", "l1pfCandidates"));
0189   descriptions.add("L1NNTauProducer", desc);
0190 }
0191 
0192 void L1NNTauProducer::makeTau_HW(const l1t::PFCandidate& seed,
0193                                  l1t::PFCandidateCollection& parts,
0194                                  std::unique_ptr<l1t::PFTauCollection>& iTaus) {
0195   // Seed Cone Jet algorithm with ap_fixed types and hardware emulation
0196   L1TauEmu::detaphi_t rCone2 =
0197       L1TauEmu::detaphi_t(fTauSize_ * fTauSize_ * L1TauEmu::etaphi_base * L1TauEmu::etaphi_base);
0198   unsigned lId = 0;
0199 
0200   input2_t p1_tot = 0;
0201   input2_t p1x_tot = 0;
0202   input2_t p1y_tot = 0;
0203   input2_t p1z_tot = 0;
0204 
0205   input_t e1ta_1 = seed.eta();
0206   input_t p1hi_1 = seed.phi();
0207   L1TauEmu::pt_t pt = 0;
0208   L1TauEmu::z0_t z0 = 0;
0209   L1TauEmu::dxy_t dxy = 0;
0210 
0211   // Reconstruct the Tau Cone
0212   for (unsigned i0 = 0; i0 < parts.size(); i0++) {
0213     if (L1TauEmu::inCone(seed, (parts[i0]), rCone2)) {
0214       if (parts[i0].id() == l1t::PFCandidate::Electron || parts[i0].id() == l1t::PFCandidate::ChargedHadron ||
0215           parts[i0].id() == l1t::PFCandidate::Photon) {
0216         lId++;
0217         pt = pt + L1TauEmu::pt_t(parts[i0].pt());
0218 
0219         input2_t d1eta = input_t(parts[i0].eta()) - e1ta_1;
0220         input2_t d1phi = input_t(parts[i0].phi()) - p1hi_1;
0221         input2_t d1r2 = d1eta * d1eta + d1phi * d1phi;
0222         input2_t tmppt = input_t(parts[i0].pt());
0223         input2_t half = 0.5;
0224         p1z_tot = p1z_tot + tmppt * (1 - d1r2 * half);
0225         p1y_tot = p1y_tot + tmppt * d1phi;
0226         p1x_tot = p1x_tot + tmppt * d1eta;
0227         p1_tot = p1_tot + tmppt;
0228 
0229         if (z0 == 0 && parts[i0].id() == l1t::PFCandidate::ChargedHadron) {
0230           z0 = parts[i0].hwZ0();
0231           dxy = parts[i0].hwDxy();
0232         }
0233       }
0234     }
0235   }
0236 
0237   //Compute the mass
0238   input2_t tmpmass1 = (p1_tot * p1_tot - p1x_tot * p1x_tot - p1y_tot * p1y_tot - p1z_tot * p1z_tot);
0239   if (tmpmass1 < 0)
0240     tmpmass1 = 0;
0241   L1TauEmu::pt_t mass = l1ct::pt_t(tmpmass1);
0242 
0243   if (pt < fSeedPt_)
0244     return;
0245 
0246   // Tau NN Inference
0247   Tau_NN_Result NN_ouput = fTauNNIdHW_->compute(seed, parts);
0248 
0249   // Needed for making PFTau
0250   input_t* lNNVector = fTauNNIdHW_->NNVectorVar();
0251   float pNNVec[80];
0252   for (unsigned i0 = 0; i0 < 80; i0++)
0253     pNNVec[i0] = float(lNNVector[i0]);
0254 
0255   //Firmware Tau
0256   l1ct::Tau l1ctTau;
0257   l1ctTau.hwPt = l1ct::pt_t(pt * NN_ouput.nn_pt_correction);  //l1gt is <16,11> and currently <16,14>
0258   l1ctTau.hwEta = l1ct::Scales::makeGlbEta(seed.eta());       // seed.eta() and seed.phi() are in physical coordinates
0259   l1ctTau.hwPhi = l1ct::Scales::makeGlbPhi(seed.phi());
0260 
0261   l1ctTau.hwSeedPt = seed.pt();
0262   l1ctTau.hwSeedZ0 = seed.hwZ0();
0263   l1ctTau.hwCharge = seed.charge();
0264 
0265   l1ctTau.hwType = l1ct::Tau::type_t(lId);
0266   l1ctTau.hwRawId = ap_uint<10>(NN_ouput.nn_id * 1024);  //NN Output is ap_fixed<16, 6> so need to cast.
0267 
0268   //Convert to GT format and pack to encodedTau of PFTau
0269   l1gt::Tau l1gtTau = l1ctTau.toGT();
0270   l1gt::PackedTau packed_Tau = l1gtTau.pack();
0271 
0272   //Make PFTau
0273   //Save pt, eta and phi in gt scales
0274   math::PtEtaPhiMLorentzVector tempP4(l1gt::Scales::floatPt(l1gtTau.v3.pt),
0275                                       l1gt::Scales::floatEta(l1gtTau.v3.eta),
0276                                       l1gt::Scales::floatPhi(l1gtTau.v3.phi),
0277                                       float(mass));
0278 
0279   l1t::PFTau l1PFTau(tempP4, pNNVec, NN_ouput.nn_id, 0, lId);
0280   l1PFTau.setZ0(float(z0) * 0.05);    //L1TauEmu::z0_base);
0281   l1PFTau.setDxy(float(dxy) * 0.05);  //L1TauEmu::dxy_base);
0282 
0283   l1PFTau.set_encodedTau(packed_Tau);
0284 
0285   iTaus->push_back(l1PFTau);
0286 }
0287 
0288 void L1NNTauProducer::process_HW(const l1t::PFCandidateCollection& parts,
0289                                  std::unique_ptr<l1t::PFTauCollection>& iTaus) {
0290   // The fixed point algorithm emulation
0291   using namespace L1TauEmu;
0292   std::vector<l1t::PFCandidate> work;
0293   work.resize(parts.size());
0294   std::transform(parts.begin(), parts.end(), work.begin(), [](const l1t::PFCandidate& part) { return part; });
0295   std::sort(work.begin(), work.end(), [](l1t::PFCandidate i, l1t::PFCandidate j) {
0296     return (l1ct::pt_t(i.pt()) > l1ct::pt_t(j.pt()));
0297   });
0298 
0299   std::vector<l1t::PFCandidate> seeds;
0300   uint lSeed = l1t::PFCandidate::ChargedHadron;
0301   if (fEMSeed)
0302     lSeed = l1t::PFCandidate::Photon;
0303   std::copy_if(work.begin(), work.end(), std::back_inserter(seeds), [&](const l1t::PFCandidate& part) {
0304     return ((part.id() == l1t::PFCandidate::Electron || part.id() == l1t::PFCandidate::ChargedHadron ||
0305              part.id() == lSeed) &&
0306             std::abs(part.eta()) < track_trigger_eta_max);
0307   });
0308   // It would be nice to transform the inputs to the etaphi_base of the FW here, as in the line below
0309   // However the phi may wrap around if the etaphi_base > 1, so don't do it...
0310   //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()));});
0311   detaphi_t rCone2 = detaphi_t(fConeSize_ * fConeSize_ * etaphi_base * etaphi_base);
0312 
0313   iTaus->reserve(fMaxTaus_);
0314   while (!seeds.empty() && iTaus->size() < unsigned(fMaxTaus_)) {
0315     // Take the first (highest pt) candidate as a seed
0316     l1t::PFCandidate seed = seeds.at(0);
0317     // Get the particles within a _coneSize of the seed
0318     std::vector<l1t::PFCandidate> particlesInCone;
0319     std::copy_if(work.begin(), work.end(), std::back_inserter(particlesInCone), [&](l1t::PFCandidate& part) {
0320       return inCone(seed, part, rCone2);
0321     });
0322     makeTau_HW(seed, particlesInCone, iTaus);
0323     // remove the clustered particles
0324     work.erase(std::remove_if(
0325                    work.begin(), work.end(), [&](const l1t::PFCandidate& part) { return inCone(seed, part, rCone2); }),
0326                work.end());
0327 
0328     seeds.erase(
0329         std::remove_if(
0330             seeds.begin(), seeds.end(), [&](const l1t::PFCandidate& part) { return inCone(seed, part, rCone2); }),
0331         seeds.end());
0332   }
0333 }
0334 
0335 L1NNTauProducer::~L1NNTauProducer() {}
0336 
0337 #include "FWCore/Framework/interface/MakerMacros.h"
0338 DEFINE_FWK_MODULE(L1NNTauProducer);