Phase2L1ParticleFlow/plugins/L1MetPfProducer.cc

0001 #include <vector>
0002 #include <string>
0003 #include <ap_int.h>
0004 #include <ap_fixed.h>
0005 #include <TVector2.h>
0006
0007 #include "FWCore/Framework/interface/global/EDProducer.h"
0008 #include "FWCore/Framework/interface/Event.h"
0009 #include "FWCore/Framework/interface/MakerMacros.h"
0010 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0011 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0012
0013 #include "DataFormats/L1TParticleFlow/interface/PFCandidate.h"
0014 #include "DataFormats/L1Trigger/interface/EtSum.h"
0015 #include "DataFormats/Math/interface/LorentzVector.h"
0016
0017 #include "hls4ml/emulator.h"
0018
0019 using namespace l1t;
0020
0021 class L1MetPfProducer : public edm::global::EDProducer<> {
0022 public:
0023   explicit L1MetPfProducer(const edm::ParameterSet&);
0024   ~L1MetPfProducer() override;
0025   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0026
0027 private:
0028   void produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
0029   edm::EDGetTokenT<vector<l1t::PFCandidate>> _l1PFToken;
0030
0031   int maxCands_ = 128;
0032
0033   // quantization controllers
0034   typedef ap_ufixed<14, 12, AP_RND, AP_WRAP> pt_t;  // LSB is 0.25 and max is 4 TeV
0035   typedef ap_int<12> phi_t;                         // LSB is pi/720 ~ 0.0044 and max is +/-8.9
0036   static constexpr float ptLSB_ = 0.25;             // GeV
0037   static constexpr float phiLSB_ = M_PI / 720;      // rad
0038
0039   // derived, helper types
0040   typedef ap_fixed<pt_t::width + 1, pt_t::iwidth + 1, AP_RND, AP_SAT> pxy_t;
0041   typedef ap_fixed<2 * pt_t::width, 2 * pt_t::iwidth, AP_RND, AP_SAT> pt2_t;
0042   // derived, helper constants
0043   static constexpr float maxPt_ = ((1 << pt_t::width) - 1) * ptLSB_;
0044   const phi_t hwPi_ = round(M_PI / phiLSB_);
0045   const phi_t hwPiOverTwo_ = round(M_PI / (2 * phiLSB_));
0046
0047   typedef ap_ufixed<pt_t::width, 0> inv_t;  // can't easily use the MAXPT/pt trick with ap_fixed
0048
0049   // to make configurable...
0050   static constexpr int dropBits_ = 2;
0051   static constexpr int dropFactor_ = (1 << dropBits_);
0052   static constexpr int invTableBits_ = 10;
0053   static constexpr int invTableSize_ = (1 << invTableBits_);
0054
0055   // hls4ml emulator objects
0056   bool useMlModel_;
0057   std::shared_ptr<hls4mlEmulator::Model> model;
0058   std::string modelVersion_;
0059   typedef ap_fixed<32, 16> input_t;
0060   typedef ap_fixed<32, 16> result_t;
0061   static constexpr int numContInputs_ = 4;
0062   static constexpr int numPxPyInputs_ = 2;
0063   static constexpr int numCatInputs_ = 2;
0064   static constexpr int numInputs_ = numContInputs_ + numPxPyInputs_ + numCatInputs_;
0065
0066   void Project(pt_t pt, phi_t phi, pxy_t& pxy, bool isX, bool debug = false) const;
0067   void PhiFromXY(pxy_t px, pxy_t py, phi_t& phi, bool debug = false) const;
0068
0069   int EncodePdgId(int pdgId) const;
0070
0071   void CalcMetHLS(const std::vector<float>& pt,
0072                   const std::vector<float>& phi,
0073                   reco::Candidate::PolarLorentzVector& metVector) const;
0074   void CalcMlMet(const std::vector<float>& pt,
0075                  const std::vector<float>& eta,
0076                  const std::vector<float>& phi,
0077                  const std::vector<float>& puppiWeight,
0078                  const std::vector<int>& pdgId,
0079                  const std::vector<int>& charge,
0080                  reco::Candidate::PolarLorentzVector& metVector) const;
0081 };
0082
0083 L1MetPfProducer::L1MetPfProducer(const edm::ParameterSet& cfg)
0084     : _l1PFToken(consumes<std::vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))),
0085       maxCands_(cfg.getParameter<int>("maxCands")),
0086       modelVersion_(cfg.getParameter<std::string>("modelVersion")) {
0087   produces<std::vector<l1t::EtSum>>();
0088   useMlModel_ = (!modelVersion_.empty());
0089   if (useMlModel_) {
0090     hls4mlEmulator::ModelLoader loader(modelVersion_);
0091     model = loader.load_model();
0092   }
0093 }
0094
0095 void L1MetPfProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0096   edm::ParameterSetDescription desc;
0097   desc.add<edm::InputTag>("L1PFObjects", edm::InputTag("L1PFProducer", "l1pfCandidates"));
0098   desc.add<int>("maxCands", 128);
0099   desc.add<std::string>("modelVersion", "");
0100   descriptions.add("L1MetPfProducer", desc);
0101 }
0102
0103 void L1MetPfProducer::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0104   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
0105   iEvent.getByToken(_l1PFToken, l1PFCandidates);
0106
0107   std::vector<float> pt;
0108   std::vector<float> eta;
0109   std::vector<float> phi;
0110   std::vector<float> puppiWeight;
0111   std::vector<int> pdgId;
0112   std::vector<int> charge;
0113
0114   for (int i = 0; i < int(l1PFCandidates->size()) && (i < maxCands_ || maxCands_ < 0); i++) {
0115     const auto& l1PFCand = l1PFCandidates->at(i);
0116     pt.push_back(l1PFCand.pt());
0117     eta.push_back(l1PFCand.eta());
0118     phi.push_back(l1PFCand.phi());
0119     puppiWeight.push_back(l1PFCand.puppiWeight());
0120     pdgId.push_back(l1PFCand.pdgId());
0121     charge.push_back(l1PFCand.charge());
0122   }
0123
0124   reco::Candidate::PolarLorentzVector metVector;
0125
0126   if (useMlModel_) {
0127     CalcMlMet(pt, eta, phi, puppiWeight, pdgId, charge, metVector);
0128   } else {
0129     CalcMetHLS(pt, phi, metVector);
0130   }
0131
0132   l1t::EtSum theMET(metVector, l1t::EtSum::EtSumType::kTotalHt, 0, 0, 0, 0);
0133
0134   auto metCollection = std::make_unique<std::vector<l1t::EtSum>>(0);
0135   metCollection->push_back(theMET);
0136   iEvent.put(std::move(metCollection));
0137 }
0138
0139 int L1MetPfProducer::EncodePdgId(int pdgId) const {
0140   switch (abs(pdgId)) {
0141     case 211:  // charged hadron (pion)
0142       return 1;
0143     case 130:  // neutral hadron (kaon)
0144       return 2;
0145     case 22:  // photon
0146       return 3;
0147     case 13:  // muon
0148       return 4;
0149     case 11:  // electron
0150       return 5;
0151     default:
0152       return 0;
0153   }
0154 }
0155
0156 void L1MetPfProducer::CalcMlMet(const std::vector<float>& pt,
0157                                 const std::vector<float>& eta,
0158                                 const std::vector<float>& phi,
0159                                 const std::vector<float>& puppiWeight,
0160                                 const std::vector<int>& pdgId,
0161                                 const std::vector<int>& charge,
0162                                 reco::Candidate::PolarLorentzVector& metVector) const {
0163   const int inputSize = maxCands_ * numInputs_;
0164
0165   input_t input[800];
0166   result_t result[2];
0167
0168   // initialize with zeros (for padding)
0169   for (int i = 0; i < inputSize; i++) {
0170     input[i] = 0;
0171   }
0172
0173   for (uint i = 0; i < pt.size(); i++) {
0174     // input_cont
0175     input[i * numContInputs_] = pt[i];
0176     input[i * numContInputs_ + 1] = eta[i];
0177     input[i * numContInputs_ + 2] = phi[i];
0178     input[i * numContInputs_ + 3] = puppiWeight[i];
0179     // input_pxpy
0180     input[(maxCands_ * numContInputs_) + (i * numPxPyInputs_)] = pt[i] * cos(phi[i]);
0181     input[(maxCands_ * numContInputs_) + (i * numPxPyInputs_) + 1] = pt[i] * sin(phi[i]);
0182     // input_cat0
0183     input[maxCands_ * (numContInputs_ + numPxPyInputs_) + i] = EncodePdgId(pdgId[i]);
0184     // input_cat1
0185     input[maxCands_ * (numContInputs_ + numPxPyInputs_ + 1) + i] = (abs(charge[i]) <= 1) ? (charge[i] + 2) : 0;
0186   }
0187
0188   model->prepare_input(input);
0189   model->predict();
0190   model->read_result(result);
0191
0192   double met_px = -result[0].to_double();
0193   double met_py = -result[1].to_double();
0194   metVector.SetPt(hypot(met_px, met_py));
0195   metVector.SetPhi(atan2(met_py, met_px));
0196   metVector.SetEta(0);
0197 }
0198
0199 void L1MetPfProducer::CalcMetHLS(const std::vector<float>& pt,
0200                                  const std::vector<float>& phi,
0201                                  reco::Candidate::PolarLorentzVector& metVector) const {
0202   pxy_t hw_px = 0;
0203   pxy_t hw_py = 0;
0204   pxy_t hw_sumx = 0;
0205   pxy_t hw_sumy = 0;
0206
0207   for (uint i = 0; i < pt.size(); i++) {
0208     pt_t hw_pt = min(pt[i], maxPt_);
0209     phi_t hw_phi = float(TVector2::Phi_mpi_pi(phi[i]) / phiLSB_);
0210
0211     Project(hw_pt, hw_phi, hw_px, true);
0212     Project(hw_pt, hw_phi, hw_py, false);
0213
0214     hw_sumx = hw_sumx - hw_px;
0215     hw_sumy = hw_sumy - hw_py;
0216   }
0217
0218   pt2_t hw_met = pt2_t(hw_sumx) * pt2_t(hw_sumx) + pt2_t(hw_sumy) * pt2_t(hw_sumy);
0219   hw_met = sqrt(int(hw_met));  // stand-in for HLS::sqrt
0220
0221   phi_t hw_met_phi = 0;
0222   PhiFromXY(hw_sumx, hw_sumy, hw_met_phi);
0223
0224   metVector.SetPt(hw_met.to_double());
0225   metVector.SetPhi(hw_met_phi.to_double() * phiLSB_);
0226   metVector.SetEta(0);
0227 }
0228
0229 void L1MetPfProducer::Project(pt_t pt, phi_t phi, pxy_t& pxy, bool isX, bool debug) const {
0230   /*
0231       Convert pt and phi to px (py)
0232       1) Map phi to the first quadrant to reduce LUT size
0233       2) Lookup sin(phiQ1), where the result is in [0,maxPt]
0234       which is used to encode [0,1].
0235       3) Multiply pt by sin(phiQ1) to get px. Result will be px*maxPt, but
0236       wrapping multiplication is 'mod maxPt' so the correct value is returned.
0237       4) Check px=-|px|.
0238     */
0239
0240   // set phi to first quadrant
0241   phi_t phiQ1 = (phi > 0) ? phi : phi_t(-phi);  // Q1/Q4
0242   if (phiQ1 >= hwPiOverTwo_)
0243     phiQ1 = hwPi_ - phiQ1;
0244
0245   if (phiQ1 > hwPiOverTwo_) {
0246     edm::LogWarning("L1MetPfProducer") << "unexpected phi (high)";
0247     phiQ1 = hwPiOverTwo_;
0248   } else if (phiQ1 < 0) {
0249     edm::LogWarning("L1MetPfProducer") << "unexpected phi (low)";
0250     phiQ1 = 0;
0251   }
0252   if (isX) {
0253     typedef ap_ufixed<14, 12, AP_RND, AP_WRAP> pt_t;  // LSB is 0.25 and max is 4 TeV
0254     ap_ufixed<pt_t::width, 0> cosPhi = cos(phiQ1.to_double() / hwPiOverTwo_.to_double() * M_PI / 2);
0255     pxy = pt * cosPhi;
0256     if (phi > hwPiOverTwo_ || phi < -hwPiOverTwo_)
0257       pxy = -pxy;
0258   } else {
0259     ap_ufixed<pt_t::width, 0> sinPhi = sin(phiQ1.to_double() / hwPiOverTwo_.to_double() * M_PI / 2);
0260     pxy = pt * sinPhi;
0261     if (phi < 0)
0262       pxy = -pxy;
0263   }
0264 }
0265
0266 void L1MetPfProducer::PhiFromXY(pxy_t px, pxy_t py, phi_t& phi, bool debug) const {
0267   if (px == 0 && py == 0) {
0268     phi = 0;
0269     return;
0270   }
0271   if (px == 0) {
0272     phi = py > 0 ? hwPiOverTwo_ : phi_t(-hwPiOverTwo_);
0273     return;
0274   }
0275   if (py == 0) {
0276     phi = px > 0 ? phi_t(0) : phi_t(-hwPi_);
0277     return;
0278   }
0279
0280   // get q1 coordinates
0281   pt_t x = px > 0 ? pt_t(px) : pt_t(-px);  //px>=0 ? px : -px;
0282   pt_t y = py > 0 ? pt_t(py) : pt_t(-py);  //px>=0 ? px : -px;
0283   // transform so a<b
0284   pt_t a = x < y ? x : y;
0285   pt_t b = x < y ? y : x;
0286
0287   if (b.to_double() > maxPt_ / dropFactor_)
0288     b = maxPt_ / dropFactor_;
0289   // map [0,max/4) to inv table size
0290   int index = round((b.to_double() / (maxPt_ / dropFactor_)) * invTableSize_);
0291   float bcheck = (float(index) / invTableSize_) * (maxPt_ / dropFactor_);
0292   inv_t inv_b = 1. / ((float(index) / invTableSize_) * (maxPt_ / dropFactor_));
0293
0294   inv_t a_over_b = a * inv_b;
0295
0296   if (debug) {
0297     printf("  a, b = %f, %f;   index, inv = %d, %f; ratio = %f \n",
0298            a.to_double(),
0299            b.to_double(),
0300            index,
0301            inv_b.to_double(),
0302            a_over_b.to_double());
0303     printf("    bcheck, 1/bc = %f, %f  -- %d  %f  %d  \n", bcheck, 1. / bcheck, invTableSize_, maxPt_, dropFactor_);
0304   }
0305
0306   int atanTableBits_ = 7;
0307   int atanTableSize_ = (1 << atanTableBits_);
0308   index = round(a_over_b.to_double() * atanTableSize_);
0309   phi = atan(float(index) / atanTableSize_) / phiLSB_;
0310
0311   if (debug) {
0312     printf("    atan index, phi = %d, %f (%f rad)  real atan(a/b)= %f  \n",
0313            index,
0314            phi.to_double(),
0315            phi.to_double() * (M_PI / hwPi_.to_double()),
0316            atan(a.to_double() / b.to_double()));
0317   }
0318
0319   // rotate from (0,pi/4) to full quad1
0320   if (y > x)
0321     phi = hwPiOverTwo_ - phi;  //phi = pi/2 - phi
0322   // other quadrants
0323   if (px < 0 && py > 0)
0324     phi = hwPi_ - phi;  // Q2 phi = pi - phi
0325   if (px > 0 && py < 0)
0326     phi = -phi;  // Q4 phi = -phi
0327   if (px < 0 && py < 0)
0328     phi = -(hwPi_ - phi);  // Q3 composition of both
0329
0330   if (debug) {
0331     printf("    phi hw, float, real = %f, %f    (%f rad from x,y = %f, %f) \n",
0332            phi.to_double(),
0333            phi.to_double() * (M_PI / hwPi_.to_double()),
0334            atan2(py.to_double(), px.to_double()),
0335            px.to_double(),
0336            py.to_double());
0337   }
0338 }
0339
0340 L1MetPfProducer::~L1MetPfProducer() {}
0341
0342 #include "FWCore/Framework/interface/MakerMacros.h"
0343 DEFINE_FWK_MODULE(L1MetPfProducer);