Phase2L1ParticleFlow/src/JetId.cc

0001 #include "L1Trigger/Phase2L1ParticleFlow/interface/JetId.h"
0002 #include "DataFormats/Math/interface/deltaPhi.h"
0003 #include <cmath>
0004
0005 JetId::JetId(const std::string &iInput,
0006              const std::string &iOutput,
0007              const std::shared_ptr<hls4mlEmulator::Model> model,
0008              int iNParticles)
0009     : modelRef_(model) {
0010   NNvectorVar_.clear();
0011   fNParticles_ = iNParticles;
0012
0013   fPt_ = std::make_unique<float[]>(fNParticles_);
0014   fEta_ = std::make_unique<float[]>(fNParticles_);
0015   fPhi_ = std::make_unique<float[]>(fNParticles_);
0016   fId_ = std::make_unique<int[]>(fNParticles_);
0017   fCharge_ = std::make_unique<int[]>(fNParticles_);
0018   fDZ_ = std::make_unique<float[]>(fNParticles_);
0019   fDX_ = std::make_unique<float[]>(fNParticles_);
0020   fDY_ = std::make_unique<float[]>(fNParticles_);
0021   fInput_ = iInput;
0022   fOutput_ = iOutput;
0023 }
0024
0025 //--BJet algo specific constructor
0026 JetId::JetId(const std::string &iInput, const std::string &iOutput, const BJetTFCache *cache, int iNParticles)
0027     : sessionRef_(cache->session) {
0028   NNvectorVar_.clear();
0029   fNParticles_ = iNParticles;
0030
0031   fPt_ = std::make_unique<float[]>(fNParticles_);
0032   fEta_ = std::make_unique<float[]>(fNParticles_);
0033   fPhi_ = std::make_unique<float[]>(fNParticles_);
0034   fId_ = std::make_unique<int[]>(fNParticles_);
0035   fCharge_ = std::make_unique<int[]>(fNParticles_);
0036   fDZ_ = std::make_unique<float[]>(fNParticles_);
0037   fDX_ = std::make_unique<float[]>(fNParticles_);
0038   fDY_ = std::make_unique<float[]>(fNParticles_);
0039   fInput_ = iInput;
0040   fOutput_ = iOutput;
0041 }
0042
0043 void JetId::setNNVectorVar() {
0044   NNvectorVar_.clear();
0045   for (int i0 = 0; i0 < fNParticles_; i0++) {
0046     if (fPt_.get()[i0] == 0) {
0047       for (int i1 = 0; i1 < 13; i1++)
0048         NNvectorVar_.push_back(0);
0049       continue;
0050     }
0051     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::Electron && fCharge_.get()[i0] < 0);       // Electron
0052     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::Electron && fCharge_.get()[i0] > 0);       // Positron
0053     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::Muon && fCharge_.get()[i0] < 0);           // Muon
0054     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::Muon && fCharge_.get()[i0] > 0);           // Anti-Muon
0055     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::Photon);                                   // Photon
0056     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::NeutralHadron);                            // Neutral Had
0057     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::ChargedHadron && fCharge_.get()[i0] < 0);  // Pion
0058     NNvectorVar_.push_back(fId_.get()[i0] == l1t::PFCandidate::ChargedHadron && fCharge_.get()[i0] > 0);  // Anti-Pion
0059     NNvectorVar_.push_back(fDZ_.get()[i0]);                                                               //dZ
0060     NNvectorVar_.push_back(std::hypot(fDX_.get()[i0], fDY_.get()[i0]));                                   //d0
0061     NNvectorVar_.push_back(fPt_.get()[i0]);   //pT as a fraction of jet pT
0062     NNvectorVar_.push_back(fEta_.get()[i0]);  //dEta from jet axis
0063     NNvectorVar_.push_back(fPhi_.get()[i0]);  //dPhi from jet axis
0064   }
0065 }
0066 float JetId::EvaluateNN() {
0067   tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, (unsigned int)NNvectorVar_.size(), 1});
0068   for (unsigned int i = 0; i < NNvectorVar_.size(); i++) {
0069     input.tensor<float, 3>()(0, i, 0) = float(NNvectorVar_[i]);
0070   }
0071   std::vector<tensorflow::Tensor> outputs;
0072   tensorflow::run(sessionRef_, {{fInput_, input}}, {fOutput_}, &outputs);
0073   return outputs[0].matrix<float>()(0, 0);
0074 }  //end EvaluateNN
0075
0076 ap_fixed<16, 6> JetId::EvaluateNNFixed() {
0077   ap_fixed<16, 6> modelInput[140] = {};
0078   for (unsigned int i = 0; i < NNvectorVar_.size(); i++) {
0079     modelInput[i] = NNvectorVar_[i];
0080   }
0081   ap_fixed<16, 6> modelResult[1] = {-1};
0082
0083   modelRef_->prepare_input(modelInput);
0084   modelRef_->predict();
0085   modelRef_->read_result(modelResult);
0086   ap_fixed<16, 6> modelResult_ = modelResult[0];
0087   return modelResult_;
0088 }  //end EvaluateNNFixed
0089
0090 float JetId::compute(const l1t::PFJet &iJet, float vz, bool useRawPt) {
0091   for (int i0 = 0; i0 < fNParticles_; i0++) {
0092     fPt_.get()[i0] = 0;
0093     fEta_.get()[i0] = 0;
0094     fPhi_.get()[i0] = 0;
0095     fId_.get()[i0] = 0;
0096     fCharge_.get()[i0] = 0;
0097     fDZ_.get()[i0] = 0;
0098     fDX_.get()[i0] = 0;
0099     fDY_.get()[i0] = 0;
0100   }
0101   auto iParts = iJet.constituents();
0102   std::sort(iParts.begin(), iParts.end(), [](edm::Ptr<l1t::PFCandidate> i, edm::Ptr<l1t::PFCandidate> j) {
0103     return (i->pt() > j->pt());
0104   });
0105   float jetpt = useRawPt ? iJet.rawPt() : iJet.pt();
0106   for (unsigned int i0 = 0; i0 < iParts.size(); i0++) {
0107     if (i0 >= (unsigned int)fNParticles_)
0108       break;
0109     fPt_.get()[i0] = iParts[i0]->pt() / jetpt;
0110     fEta_.get()[i0] = iParts[i0]->eta() - iJet.eta();
0111     fPhi_.get()[i0] = deltaPhi(iParts[i0]->phi(), iJet.phi());
0112     fId_.get()[i0] = iParts[i0]->id();
0113     fCharge_.get()[i0] = iParts[i0]->charge();
0114     if (iParts[i0]->pfTrack().isNonnull()) {
0115       fDX_.get()[i0] = iParts[i0]->pfTrack()->vx();
0116       fDY_.get()[i0] = iParts[i0]->pfTrack()->vy();
0117       fDZ_.get()[i0] = iParts[i0]->pfTrack()->vz() - vz;
0118     }
0119   }
0120   setNNVectorVar();
0121   return EvaluateNN();
0122 }
0123
0124 ap_fixed<16, 6> JetId::computeFixed(const l1t::PFJet &iJet, float vz, bool useRawPt) {
0125   for (int i0 = 0; i0 < fNParticles_; i0++) {
0126     fPt_.get()[i0] = 0;
0127     fEta_.get()[i0] = 0;
0128     fPhi_.get()[i0] = 0;
0129     fId_.get()[i0] = 0;
0130     fCharge_.get()[i0] = 0;
0131     fDZ_.get()[i0] = 0;
0132     fDX_.get()[i0] = 0;
0133     fDY_.get()[i0] = 0;
0134   }
0135   auto iParts = iJet.constituents();
0136   std::sort(iParts.begin(), iParts.end(), [](edm::Ptr<l1t::PFCandidate> i, edm::Ptr<l1t::PFCandidate> j) {
0137     return (i->pt() > j->pt());
0138   });
0139   float jetpt = useRawPt ? iJet.rawPt() : iJet.pt();
0140   for (unsigned int i0 = 0; i0 < iParts.size(); i0++) {
0141     if (i0 >= (unsigned int)fNParticles_)
0142       break;
0143     fPt_.get()[i0] = iParts[i0]->pt() / jetpt;
0144     fEta_.get()[i0] = iParts[i0]->eta() - iJet.eta();
0145     fPhi_.get()[i0] = deltaPhi(iParts[i0]->phi(), iJet.phi());
0146     fId_.get()[i0] = iParts[i0]->id();
0147     fCharge_.get()[i0] = iParts[i0]->charge();
0148     if (iParts[i0]->pfTrack().isNonnull()) {
0149       fDX_.get()[i0] = iParts[i0]->pfTrack()->vx();
0150       fDY_.get()[i0] = iParts[i0]->pfTrack()->vy();
0151       fDZ_.get()[i0] = iParts[i0]->pfTrack()->vz() - vz;
0152     }
0153   }
0154   setNNVectorVar();
0155   return EvaluateNNFixed();
0156 }