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
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 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
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
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
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
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
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
0247 Tau_NN_Result NN_ouput = fTauNNIdHW_->compute(seed, parts);
0248
0249
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
0256 l1ct::Tau l1ctTau;
0257 l1ctTau.hwPt = l1ct::pt_t(pt * NN_ouput.nn_pt_correction);
0258 l1ctTau.hwEta = l1ct::Scales::makeGlbEta(seed.eta());
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);
0267
0268
0269 l1gt::Tau l1gtTau = l1ctTau.toGT();
0270 l1gt::PackedTau packed_Tau = l1gtTau.pack();
0271
0272
0273
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);
0281 l1PFTau.setDxy(float(dxy) * 0.05);
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
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
0309
0310
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
0316 l1t::PFCandidate seed = seeds.at(0);
0317
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
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);