Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-11 23:28:04

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