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
0033 std::unique_ptr<TauNNId> fTauNNId_;
0034 std::unique_ptr<TauNNIdHW> fTauNNIdHW_;
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");
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
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
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
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
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
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
0248 Tau_NN_Result NN_ouput = fTauNNIdHW_->compute(seed, parts);
0249
0250
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
0257 l1ct::Tau l1ctTau;
0258 l1ctTau.hwPt = l1ct::pt_t(pt * NN_ouput.nn_pt_correction);
0259 l1ctTau.hwEta = l1ct::Scales::makeGlbEta(seed.eta());
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);
0268
0269
0270 l1gt::Tau l1gtTau = l1ctTau.toGT();
0271 l1gt::PackedTau packed_Tau = l1gtTau.pack();
0272
0273
0274
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);
0282 l1PFTau.setDxy(float(dxy) * 0.05);
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
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
0310
0311
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
0317 l1t::PFCandidate seed = seeds.at(0);
0318
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
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);