Add Phase-2 Level-1 trigger ML MET emulator

L1Trigger/Configuration/python/L1Trigger_EventContent_cff.py

@@ -218,6 +218,7 @@ def _appendPhase2Digis(obj):
         'keep *_l1tLayer1EG_*_*',
         'keep *_l1tLayer2EG_*_*',
         'keep *_l1tMETPFProducer_*_*',
+        'keep *_l1tMETMLProducer_*_*',
         'keep *_l1tNNTauProducer_*_*',
         'keep *_l1tNNTauProducerPuppi_*_*',
         'keep *_l1tHPSPFTauProducerPF_*_*',

L1Trigger/Configuration/python/SimL1Emulator_cff.py

@@ -223,6 +223,7 @@
 
 from L1Trigger.Phase2L1ParticleFlow.l1tMETPFProducer_cfi import *
 _phase2_siml1emulator.add(l1tMETPFProducer)
+_phase2_siml1emulator.add(l1tMETMLProducer)
 
 
 # NNTaus

L1Trigger/Phase2L1ParticleFlow/BuildFile.xml

@@ -13,6 +13,8 @@
 <use name="tensorflow"/>
 <use name="roottmva"/>
 <use name="hls"/>
+<use name="hls4mlEmulatorExtras"/>
+<use name="L1METML"/>
 <export>
   <lib name="1"/>
 </export>

L1Trigger/Phase2L1ParticleFlow/plugins/L1MetPfProducer.cc

@@ -1,4 +1,5 @@
 #include <vector>
+#include <string>
 #include <ap_int.h>
 #include <ap_fixed.h>
 #include <TVector2.h>
@@ -13,12 +14,15 @@
 #include "DataFormats/L1Trigger/interface/EtSum.h"
 #include "DataFormats/Math/interface/LorentzVector.h"
 
+#include "hls4ml/emulator.h"
+
 using namespace l1t;
 
 class L1MetPfProducer : public edm::global::EDProducer<> {
 public:
   explicit L1MetPfProducer(const edm::ParameterSet&);
   ~L1MetPfProducer() override;
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
 
 private:
   void produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
@@ -29,63 +33,171 @@ class L1MetPfProducer : public edm::global::EDProducer<> {
   // quantization controllers
   typedef ap_ufixed<14, 12, AP_RND, AP_WRAP> pt_t;  // LSB is 0.25 and max is 4 TeV
   typedef ap_int<12> phi_t;                         // LSB is pi/720 ~ 0.0044 and max is +/-8.9
-  const float ptLSB_ = 0.25;                        // GeV
-  const float phiLSB_ = M_PI / 720;                 // rad
+  static constexpr float ptLSB_ = 0.25;             // GeV
+  static constexpr float phiLSB_ = M_PI / 720;      // rad
 
   // derived, helper types
   typedef ap_fixed<pt_t::width + 1, pt_t::iwidth + 1, AP_RND, AP_SAT> pxy_t;
   typedef ap_fixed<2 * pt_t::width, 2 * pt_t::iwidth, AP_RND, AP_SAT> pt2_t;
   // derived, helper constants
-  const float maxPt_ = ((1 << pt_t::width) - 1) * ptLSB_;
+  static constexpr float maxPt_ = ((1 << pt_t::width) - 1) * ptLSB_;
   const phi_t hwPi_ = round(M_PI / phiLSB_);
   const phi_t hwPiOverTwo_ = round(M_PI / (2 * phiLSB_));
 
   typedef ap_ufixed<pt_t::width, 0> inv_t;  // can't easily use the MAXPT/pt trick with ap_fixed
 
   // to make configurable...
-  const int dropBits_ = 2;
-  const int dropFactor_ = (1 << dropBits_);
-  const int invTableBits_ = 10;
-  const int invTableSize_ = (1 << invTableBits_);
+  static constexpr int dropBits_ = 2;
+  static constexpr int dropFactor_ = (1 << dropBits_);
+  static constexpr int invTableBits_ = 10;
+  static constexpr int invTableSize_ = (1 << invTableBits_);
+
+  // hls4ml emulator objects
+  bool useMlModel_;
+  std::shared_ptr<hls4mlEmulator::Model> model;
+  std::string modelVersion_;
+  typedef ap_fixed<32, 16> input_t;
+  typedef ap_fixed<32, 16> result_t;
+  static constexpr int numContInputs_ = 4;
+  static constexpr int numPxPyInputs_ = 2;
+  static constexpr int numCatInputs_ = 2;
+  static constexpr int numInputs_ = numContInputs_ + numPxPyInputs_ + numCatInputs_;
 
   void Project(pt_t pt, phi_t phi, pxy_t& pxy, bool isX, bool debug = false) const;
   void PhiFromXY(pxy_t px, pxy_t py, phi_t& phi, bool debug = false) const;
 
-  void CalcMetHLS(std::vector<float> pt, std::vector<float> phi, reco::Candidate::PolarLorentzVector& metVector) const;
+  int EncodePdgId(int pdgId) const;
+
+  void CalcMetHLS(const std::vector<float>& pt,
+                  const std::vector<float>& phi,
+                  reco::Candidate::PolarLorentzVector& metVector) const;
+  void CalcMlMet(const std::vector<float>& pt,
+                 const std::vector<float>& eta,
+                 const std::vector<float>& phi,
+                 const std::vector<float>& puppiWeight,
+                 const std::vector<int>& pdgId,
+                 const std::vector<int>& charge,
+                 reco::Candidate::PolarLorentzVector& metVector) const;
 };
 
 L1MetPfProducer::L1MetPfProducer(const edm::ParameterSet& cfg)
     : _l1PFToken(consumes<std::vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))),
-      maxCands_(cfg.getParameter<int>("maxCands")) {
+      maxCands_(cfg.getParameter<int>("maxCands")),
+      modelVersion_(cfg.getParameter<std::string>("modelVersion")) {
   produces<std::vector<l1t::EtSum>>();
+  useMlModel_ = (modelVersion_.length() > 0);
+  if (useMlModel_) {
+    hls4mlEmulator::ModelLoader loader(modelVersion_);
+    model = loader.load_model();
+  }
+}
+
+void L1MetPfProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("L1PFObjects", edm::InputTag("L1PFProducer", "l1pfCandidates"));
+  desc.add<int>("maxCands", 128);
+  desc.add<std::string>("modelVersion", "");
+  descriptions.add("L1MetPfProducer", desc);
 }
 
 void L1MetPfProducer::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
   iEvent.getByToken(_l1PFToken, l1PFCandidates);
 
   std::vector<float> pt;
+  std::vector<float> eta;
   std::vector<float> phi;
+  std::vector<float> puppiWeight;
+  std::vector<int> pdgId;
+  std::vector<int> charge;
 
   for (int i = 0; i < int(l1PFCandidates->size()) && (i < maxCands_ || maxCands_ < 0); i++) {
     const auto& l1PFCand = l1PFCandidates->at(i);
     pt.push_back(l1PFCand.pt());
+    eta.push_back(l1PFCand.eta());
     phi.push_back(l1PFCand.phi());
+    puppiWeight.push_back(l1PFCand.puppiWeight());
+    pdgId.push_back(l1PFCand.pdgId());
+    charge.push_back(l1PFCand.charge());
   }
 
   reco::Candidate::PolarLorentzVector metVector;
 
-  CalcMetHLS(pt, phi, metVector);
+  if (useMlModel_) {
+    CalcMlMet(pt, eta, phi, puppiWeight, pdgId, charge, metVector);
+  } else {
+    CalcMetHLS(pt, phi, metVector);
+  }
 
   l1t::EtSum theMET(metVector, l1t::EtSum::EtSumType::kTotalHt, 0, 0, 0, 0);
 
-  std::unique_ptr<std::vector<l1t::EtSum>> metCollection(new std::vector<l1t::EtSum>(0));
+  auto metCollection = std::make_unique<std::vector<l1t::EtSum>>(0);
   metCollection->push_back(theMET);
   iEvent.put(std::move(metCollection));
 }
 
-void L1MetPfProducer::CalcMetHLS(std::vector<float> pt,
-                                 std::vector<float> phi,
+int L1MetPfProducer::EncodePdgId(int pdgId) const {
+  switch (abs(pdgId)) {
+    case 211:  // charged hadron (pion)
+      return 1;
+    case 130:  // neutral hadron (kaon)
+      return 2;
+    case 22:  // photon
+      return 3;
+    case 13:  // muon
+      return 4;
+    case 11:  // electron
+      return 5;
+    default:
+      return 0;
+  }
+}
+
+void L1MetPfProducer::CalcMlMet(const std::vector<float>& pt,
+                                const std::vector<float>& eta,
+                                const std::vector<float>& phi,
+                                const std::vector<float>& puppiWeight,
+                                const std::vector<int>& pdgId,
+                                const std::vector<int>& charge,
+                                reco::Candidate::PolarLorentzVector& metVector) const {
+  const int inputSize = maxCands_ * numInputs_;
+
+  input_t input[800];
+  result_t result[2];
+
+  // initialize with zeros (for padding)
+  for (int i = 0; i < inputSize; i++) {
+    input[i] = 0;
+  }
+
+  for (uint i = 0; i < pt.size(); i++) {
+    // input_cont
+    input[i * numContInputs_] = pt[i];
+    input[i * numContInputs_ + 1] = eta[i];
+    input[i * numContInputs_ + 2] = phi[i];
+    input[i * numContInputs_ + 3] = puppiWeight[i];
+    // input_pxpy
+    input[(maxCands_ * numContInputs_) + (i * numPxPyInputs_)] = pt[i] * cos(phi[i]);
+    input[(maxCands_ * numContInputs_) + (i * numPxPyInputs_) + 1] = pt[i] * sin(phi[i]);
+    // input_cat0
+    input[maxCands_ * (numContInputs_ + numPxPyInputs_) + i] = EncodePdgId(pdgId[i]);
+    // input_cat1
+    input[maxCands_ * (numContInputs_ + numPxPyInputs_ + 1) + i] = (abs(charge[i]) <= 1) ? (charge[i] + 2) : 0;
+  }
+
+  model->prepare_input(input);
+  model->predict();
+  model->read_result(result);
+
+  double met_px = -result[0].to_double();
+  double met_py = -result[1].to_double();
+  metVector.SetPt(hypot(met_px, met_py));
+  metVector.SetPhi(atan2(met_py, met_px));
+  metVector.SetEta(0);
+}
+
+void L1MetPfProducer::CalcMetHLS(const std::vector<float>& pt,
+                                 const std::vector<float>& phi,
                                  reco::Candidate::PolarLorentzVector& metVector) const {
   pxy_t hw_px = 0;
   pxy_t hw_py = 0;

L1Trigger/Phase2L1ParticleFlow/python/l1tMETPFProducer_cfi.py

@@ -3,4 +3,11 @@
 l1tMETPFProducer = cms.EDProducer("L1MetPfProducer",
                                  L1PFObjects = cms.InputTag("l1tLayer1","Puppi"),
                                  maxCands = cms.int32(128),
+                                 modelVersion = cms.string(""),
+)
+
+l1tMETMLProducer = cms.EDProducer("L1MetPfProducer",
+                                 L1PFObjects = cms.InputTag("l1tLayer1","Puppi"),
+                                 maxCands = cms.int32(100),
+                                 modelVersion = cms.string("L1METML_v1"),
 )