Merge pull request #48179 from BenjaminRS/NNVtxAssoc_Masterv1

[L1T] E2E NNVtx Track to Vertex Association in Puppi (Phase 2)

Author: cmsbuild

L1Trigger/L1TTrackMatch/plugins/L1TrackVertexAssociationProducer.cc

@@ -204,7 +204,7 @@ class L1TrackVertexAssociationProducer : public edm::global::EDProducer<> {
       inputAssoc.tensor<float, 2>()(0, 3) = dZEmulation_rescale.to_double();
 
       // Run Association Network:
-      tensorflow::run(AssociationSesh_, {{"assoc:0", inputAssoc}}, {"Identity:0"}, &outputAssoc);
+      tensorflow::run(AssociationSesh_, {{"NNvtx_track_association:0", inputAssoc}}, {"Identity:0"}, &outputAssoc);
 
       double NNOutput = (double)outputAssoc[0].tensor<float, 2>()(0, 0);
 

L1Trigger/Phase2L1ParticleFlow/interface/NNVtxAssoc.h

@@ -0,0 +1,45 @@
+#ifndef L1TRIGGER_PHASE2L1PARTICLEFLOWS_NNVtxAssoc_H
+#define L1TRIGGER_PHASE2L1PARTICLEFLOWS_NNVtxAssoc_H
+
+// Package:      L1Trigger/Phase2L1ParticleFlow
+// Class:        NNVtxAssoc
+// Description:  Designed to run the track to vertex associations created by the E2E NNVtx.
+//               TTTrackNetworkSelector either accepts or rejects that a PF object's (t) track is associated to a vertex (v).
+// Authors:      Kai Hong Law and Benjamin Radburn-Smith
+// Created:      February 2025
+
+#include <string>
+#include "DataFormats/L1TParticleFlow/interface/layer1_emulator.h"
+
+namespace edm {
+  class ParameterSet;
+  class ParameterSetDescription;
+}  // namespace edm
+namespace tensorflow {
+  class Session;
+}  // namespace tensorflow
+using namespace l1ct;
+
+class NNVtxAssoc {
+public:
+  NNVtxAssoc(std::string AssociationGraphPath,
+             const double AssociationThreshold,
+             const std::vector<double>& AssociationNetworkZ0binning,
+             const std::vector<double>& AssociationNetworkEtaBounds,
+             const std::vector<double>& AssociationNetworkZ0ResBins);
+
+  void NNVtxAssocDebug();
+  static edm::ParameterSetDescription getParameterSetDescription();
+
+  template <typename T>
+  bool TTTrackNetworkSelector(const PFRegionEmu& region, const T& t, const l1ct::PVObjEmu& v);
+
+private:
+  tensorflow::Session* associationSesh_;
+  double associationThreshold_;
+  std::vector<double> z0_binning_;
+  std::vector<double> eta_bins_;
+  std::vector<double> res_bins_;
+  std::stringstream log_;
+};
+#endif

L1Trigger/Phase2L1ParticleFlow/interface/puppi/linpuppi_ref.h

@@ -5,13 +5,20 @@
 #include "linpuppi_bits.h"
 
 #include <vector>
+#include "L1Trigger/Phase2L1ParticleFlow/interface/NNVtxAssoc.h"
+#include <memory>
 
 namespace edm {
   class ParameterSet;
   class ParameterSetDescription;
 }  // namespace edm
 
 namespace l1ct {
+#ifdef CMSSW_GIT_HASH
+  const bool withinCMSSW_ = true;
+#else
+  const bool withinCMSSW_ = false;
+#endif
 
   class LinPuppiEmulator {
   public:
@@ -35,6 +42,12 @@ namespace l1ct {
                      double priorNe,
                      double priorPh,
                      pt_t ptCut,
+                     bool useMLAssociation = false,
+                     const double associationThreshold = 0.0,
+                     std::string associationGraphPath = "",
+                     std::vector<double> associationNetworkZ0binning = {},
+                     std::vector<double> associationNetworkEtaBounds = {},
+                     std::vector<double> associationNetworkZ0ResBins = {},
                      unsigned int nFinalSort = 0,
                      SortAlgo finalSortAlgo = SortAlgo::Insertion)
         : nTrack_(nTrack),
@@ -56,10 +69,19 @@ namespace l1ct {
           priorNe_(1, priorNe),
           priorPh_(1, priorPh),
           ptCut_(1, ptCut),
+          useMLAssociation_(useMLAssociation),
           nFinalSort_(nFinalSort ? nFinalSort : nOut),
           finalSortAlgo_(finalSortAlgo),
           debug_(false),
-          fakePuppi_(false) {}
+          fakePuppi_(false) {
+      if (useMLAssociation_ and withinCMSSW_) {
+        nnVtxAssoc_ = std::make_unique<NNVtxAssoc>(NNVtxAssoc(associationGraphPath,
+                                                              associationThreshold,
+                                                              associationNetworkZ0binning,
+                                                              associationNetworkEtaBounds,
+                                                              associationNetworkZ0ResBins));
+      }
+    }
 
     LinPuppiEmulator(unsigned int nTrack,
                      unsigned int nIn,
@@ -90,6 +112,12 @@ namespace l1ct {
                      double priorPh_1,
                      pt_t ptCut_0,
                      pt_t ptCut_1,
+                     bool useMLAssociation = false,
+                     const double associationThreshold = 0.0,
+                     std::string associationGraphPath = "",
+                     std::vector<double> associationNetworkZ0binning = {},
+                     std::vector<double> associationNetworkEtaBounds = {},
+                     std::vector<double> associationNetworkZ0ResBins = {},
                      unsigned int nFinalSort = 0,
                      SortAlgo finalSortAlgo = SortAlgo::Insertion);
 
@@ -112,6 +140,12 @@ namespace l1ct {
                      const std::vector<double> &priorNe,
                      const std::vector<double> &priorPh,
                      const std::vector<pt_t> &ptCut,
+                     bool useMLAssociation,
+                     const double associationThreshold,
+                     std::string associationGraphPath,
+                     std::vector<double> associationNetworkZ0binning,
+                     std::vector<double> associationNetworkEtaBounds,
+                     std::vector<double> associationNetworkZ0ResBins,
                      unsigned int nFinalSort,
                      SortAlgo finalSortAlgo)
         : nTrack_(nTrack),
@@ -133,6 +167,7 @@ namespace l1ct {
           priorNe_(priorNe),
           priorPh_(priorPh),
           ptCut_(ptCut),
+          useMLAssociation_(useMLAssociation),
           nFinalSort_(nFinalSort),
           finalSortAlgo_(finalSortAlgo),
           debug_(false),
@@ -211,6 +246,11 @@ namespace l1ct {
     std::vector<double> alphaSlope_, alphaZero_, alphaCrop_;
     std::vector<double> priorNe_, priorPh_;
     std::vector<pt_t> ptCut_;
+
+    // NNVtx Association:
+    bool useMLAssociation_;
+    std::unique_ptr<NNVtxAssoc> nnVtxAssoc_;
+
     unsigned int nFinalSort_;  // output after a full sort of charged + neutral
     SortAlgo finalSortAlgo_;
 

L1Trigger/Phase2L1ParticleFlow/python/l1ctLayer1_cff.py

@@ -9,6 +9,9 @@
 
 from L1Trigger.Phase2L1ParticleFlow.l1TkEgAlgoEmulator_cfi import tkEgAlgoParameters,tkEgSorterParameters
 
+from L1Trigger.Phase2L1ParticleFlow.mlAssociation_cfi import NNVtxAssociationPSet
+switchOnNNAssoc = cms.bool(False)
+
 l1tLayer1Barrel = cms.EDProducer("L1TCorrelatorLayer1Producer",
     tracks = cms.InputTag('l1tPFTracksFromL1Tracks'),
     muons = cms.InputTag('l1tSAMuonsGmt','prompt'),
@@ -81,6 +84,8 @@
         alphaCrop         = cms.vdouble(  4  ), # max. absolute value for alpha term
         priors             = cms.vdouble( 5.0 ),
         priorsPhoton      = cms.vdouble( 1.0 ),
+        useMLAssociation = switchOnNNAssoc, #Enable Association Network
+        NNVtxAssociation = NNVtxAssociationPSet,
     ),
     tkEgAlgoParameters = tkEgAlgoParameters.clone(
         nTRACK = 25,
@@ -221,6 +226,8 @@
         alphaCrop         = cms.vdouble(  3 ,  3  ), # max. absolute value for alpha term
         priors             = cms.vdouble( 5.0, 5.0 ),
         priorsPhoton      = cms.vdouble( 1.5, 1.5 ), #When e/g ID not applied, use: cms.vdouble( 3.5, 3.5, 7.0 ),
+        useMLAssociation = switchOnNNAssoc, #Enable Association Network
+        NNVtxAssociation = NNVtxAssociationPSet,
     ),
     tkEgAlgoParameters = tkEgAlgoParameters.clone(
         nTRACK = 30,
@@ -327,6 +334,8 @@
         alphaCrop         = cms.vdouble(  4  ), # max. absolute value for alpha term
         priors             = cms.vdouble( 7.0 ),
         priorsPhoton      = cms.vdouble( 5.0 ), #When e/g ID not applied, use: cms.vdouble( 3.5, 3.5, 7.0 ),
+        useMLAssociation = switchOnNNAssoc, #Enable Association Network
+        NNVtxAssociation = NNVtxAssociationPSet,
     ),
     tkEgAlgoParameters = tkEgAlgoParameters.clone(
         nTRACK = 30,
@@ -396,7 +405,9 @@
         alphaCrop         = cms.vdouble(   4   ),
         priors             = cms.vdouble(  6.0  ),
         priorsPhoton      = cms.vdouble(  6.0  ),
-        debug = cms.untracked.bool(False)
+        debug = cms.untracked.bool(False),
+        useMLAssociation = switchOnNNAssoc, #Enable Association Network
+        NNVtxAssociation = NNVtxAssociationPSet,
     ),
     tkEgAlgoParameters = tkEgAlgoParameters.clone(
         nTRACK = 5,           # to be defined

L1Trigger/Phase2L1ParticleFlow/python/mlAssociation_cfi.py

@@ -0,0 +1,11 @@
+import FWCore.ParameterSet.Config as cms
+from L1Trigger.VertexFinder.l1tVertexProducer_cfi import l1tVertexProducer
+
+# NNVertex Association Variables
+NNVtxAssociationPSet = cms.PSet(
+    associationThreshold = cms.double(0.1), #Association Network threshold for PV tracks
+    associationGraph = cms.string("L1Trigger/L1TTrackMatch/data/NNVtx_AssociationModelGraph.pb"),
+    associationNetworkZ0binning = l1tVertexProducer.VertexReconstruction.FH_HistogramParameters, #Z0 binning used for setting the input feature digitisation
+    associationNetworkEtaBounds = cms.vdouble(0.0, 0.01984126984126984, 0.03968253968253968, 0.05952380952380952, 0.07936507936507936, 0.0992063492063492, 0.11904761904761904, 0.1388888888888889, 0.15873015873015872, 0.17857142857142855, 0.1984126984126984, 0.21825396825396826, 0.23809523809523808, 0.2579365079365079, 0.2777777777777778, 0.2976190476190476, 0.31746031746031744, 0.33730158730158727, 0.3571428571428571, 0.376984126984127, 0.3968253968253968, 0.41666666666666663, 0.4365079365079365, 0.45634920634920634, 0.47619047619047616, 0.496031746031746, 0.5158730158730158, 0.5357142857142857, 0.5555555555555556, 0.5753968253968254, 0.5952380952380952, 0.615079365079365, 0.6349206349206349, 0.6547619047619048, 0.6746031746031745, 0.6944444444444444, 0.7142857142857142, 0.7341269841269841, 0.753968253968254, 0.7738095238095237, 0.7936507936507936, 0.8134920634920635, 0.8333333333333333, 0.8531746031746031, 0.873015873015873, 0.8928571428571428, 0.9126984126984127, 0.9325396825396824, 0.9523809523809523, 0.9722222222222222, 0.992063492063492, 1.0119047619047619, 1.0317460317460316, 1.0515873015873016, 1.0714285714285714, 1.0912698412698412, 1.1111111111111112, 1.130952380952381, 1.1507936507936507, 1.1706349206349205, 1.1904761904761905, 1.2103174603174602, 1.23015873015873, 1.25, 1.2698412698412698, 1.2896825396825395, 1.3095238095238095, 1.3293650793650793, 1.349206349206349, 1.369047619047619, 1.3888888888888888, 1.4087301587301586, 1.4285714285714284, 1.4484126984126984, 1.4682539682539681, 1.488095238095238, 1.507936507936508, 1.5277777777777777, 1.5476190476190474, 1.5674603174603174, 1.5873015873015872, 1.607142857142857, 1.626984126984127, 1.6468253968253967, 1.6666666666666665, 1.6865079365079365, 1.7063492063492063, 1.726190476190476, 1.746031746031746, 1.7658730158730158, 1.7857142857142856, 1.8055555555555554, 1.8253968253968254, 1.8452380952380951, 1.865079365079365, 1.8849206349206349, 1.9047619047619047, 1.9246031746031744, 1.9444444444444444, 1.9642857142857142, 1.984126984126984, 2.003968253968254, 2.0238095238095237, 2.0436507936507935, 2.0634920634920633, 2.083333333333333, 2.1031746031746033, 2.123015873015873, 2.142857142857143, 2.1626984126984126, 2.1825396825396823, 2.202380952380952, 2.2222222222222223, 2.242063492063492, 2.261904761904762, 2.2817460317460316, 2.3015873015873014, 2.321428571428571, 2.341269841269841, 2.361111111111111, 2.380952380952381, 2.4007936507936507, 2.4206349206349205, 2.4404761904761902, 2.46031746031746, 2.4801587301587302, 2.5), #Eta bounds used to set z0 resolution input feature
+    associationNetworkZ0ResBins = cms.vdouble(127.0, 126.0, 126.0, 126.0, 125.0, 124.0, 123.0, 122.0, 120.0, 119.0, 117.0, 115.0, 114.0, 112.0, 110.0, 107.0, 105.0, 103.0, 101.0, 98.0, 96.0, 94.0, 91.0, 89.0, 87.0, 85.0, 82.0, 80.0, 78.0, 76.0, 74.0, 72.0, 70.0, 68.0, 66.0, 64.0, 62.0, 61.0, 59.0, 57.0, 56.0, 54.0, 53.0, 51.0, 50.0, 48.0, 47.0, 46.0, 45.0, 43.0, 42.0, 41.0, 40.0, 39.0, 38.0, 37.0, 36.0, 35.0, 34.0, 33.0, 33.0, 32.0, 31.0, 30.0, 30.0, 29.0, 28.0, 28.0, 27.0, 26.0, 26.0, 25.0, 24.0, 24.0, 23.0, 23.0, 22.0, 22.0, 21.0, 21.0, 21.0, 20.0, 20.0, 19.0, 19.0, 18.0, 18.0, 18.0, 17.0, 17.0, 17.0, 16.0, 16.0, 16.0, 15.0, 15.0, 15.0, 15.0, 14.0, 14.0, 14.0, 14.0, 13.0, 13.0, 13.0, 13.0, 12.0, 12.0, 12.0, 12.0, 12.0, 11.0, 11.0, 11.0, 11.0, 11.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 9.0, 9.0, 9.0, 9.0, 9.0, 0.0), #z0 resolution input feature bins
+)

L1Trigger/Phase2L1ParticleFlow/src/NNVtxAssoc.cc

@@ -0,0 +1,108 @@
+// Package:      L1Trigger/Phase2L1ParticleFlow
+// Class:        NNVtxAssoc
+// Description:  Designed to run the track to vertex associations created by the E2E NNVtx.
+//               TTTrackNetworkSelector either accepts or rejects that a PF object's (t) track is associated to a vertex (v).
+// Authors:      Kai Hong Law and Benjamin Radburn-Smith
+// Created:      February 2025
+
+#include "L1Trigger/Phase2L1ParticleFlow/interface/NNVtxAssoc.h"
+#include "DataFormats/L1TParticleFlow/interface/PFTrack.h"
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include <iomanip>
+
+NNVtxAssoc::NNVtxAssoc(std::string AssociationGraphPath,
+                       const double AssociationThreshold,
+                       const std::vector<double>& AssociationNetworkZ0binning,
+                       const std::vector<double>& AssociationNetworkEtaBounds,
+                       const std::vector<double>& AssociationNetworkZ0ResBins)
+    : associationThreshold_(AssociationThreshold),
+      z0_binning_(AssociationNetworkZ0binning),
+      eta_bins_(AssociationNetworkEtaBounds),
+      res_bins_(AssociationNetworkZ0ResBins) {
+  tensorflow::GraphDef* associationGraph_ = tensorflow::loadGraphDef(AssociationGraphPath);
+  associationSesh_ = tensorflow::createSession(associationGraph_);
+  log_.setf(std::ios::fixed, std::ios::floatfield);
+  log_.precision(3);
+}
+
+template <typename T>
+bool NNVtxAssoc::TTTrackNetworkSelector(const PFRegionEmu& region, const T& t, const l1ct::PVObjEmu& v) {
+  tensorflow::Tensor inputAssoc(tensorflow::DT_FLOAT, {1, 4});
+  std::vector<tensorflow::Tensor> outputAssoc;
+
+  auto lower = std::lower_bound(eta_bins_.begin(), eta_bins_.end(), region.floatGlbEta(t.hwVtxEta()));
+
+  int resbin = std::distance(eta_bins_.begin(), lower);
+  float binWidth = z0_binning_[2];
+  // Calculate integer dZ from track z0 and vertex z0 (use floating point version and convert internally allowing use of both emulator and simulator vertex and track)
+  float dZ =
+      abs(floor(((t.floatZ0() + z0_binning_[1]) / (binWidth))) - floor(((v.floatZ0() + z0_binning_[1]) / (binWidth))));
+
+  // The following constants <22, 9> are defined by the quantisation of the Neural Network
+  ap_ufixed<22, 9> ptEmulation_rescale = t.hwPt;
+  ap_ufixed<22, 9> resBinEmulation_rescale = res_bins_[resbin];
+  ap_ufixed<22, 9> MVAEmulation_rescale = 0;
+  ap_ufixed<22, 9> dZEmulation_rescale = dZ;
+
+  // Deal with this template class using 2 different objects (t) which have different calls to their PFTracks:
+  const l1t::PFTrack* srcTrack = nullptr;
+  if constexpr (std::is_same_v<T, const l1ct::TkObjEmu>)
+    srcTrack = t.src;
+  else if constexpr (std::is_same_v<T, const l1ct::PFChargedObjEmu>)
+    srcTrack = t.srcTrack;
+  if (srcTrack)
+    MVAEmulation_rescale = srcTrack->trackWord().getMVAQualityBits();
+
+  inputAssoc.tensor<float, 2>()(0, 0) = ptEmulation_rescale.to_double();
+  inputAssoc.tensor<float, 2>()(0, 1) = MVAEmulation_rescale.to_double();
+  inputAssoc.tensor<float, 2>()(0, 2) = resBinEmulation_rescale.to_double() / 16.0;
+  inputAssoc.tensor<float, 2>()(0, 3) = dZEmulation_rescale.to_double();
+
+  // Run Association Network:
+  tensorflow::run(associationSesh_, {{"NNvtx_track_association:0", inputAssoc}}, {"Identity:0"}, &outputAssoc);
+
+  double NNOutput = (double)outputAssoc[0].tensor<float, 2>()(0, 0);
+  double NNOutput_exp = 1.0 / (1.0 + exp(-1.0 * (NNOutput)));
+
+  return NNOutput_exp >= associationThreshold_;
+}
+
+#ifdef CMSSW_GIT_HASH
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+
+edm::ParameterSetDescription NNVtxAssoc::getParameterSetDescription() {
+  edm::ParameterSetDescription description;
+  description.add<double>("associationThreshold");
+  description.add<std::string>("associationGraph");
+  description.add<std::vector<double>>("associationNetworkZ0binning");
+  description.add<std::vector<double>>("associationNetworkEtaBounds");
+  description.add<std::vector<double>>("associationNetworkZ0ResBins");
+  return description;
+}
+#endif
+
+void NNVtxAssoc::NNVtxAssocDebug() {
+  log_ << "-- NNVtxAssocDebug --\n";
+  log_ << "AssociationThreshold: " << this->associationThreshold_ << "\n";
+  log_ << "z0_binning: ";
+  for (auto i : this->z0_binning_)
+    log_ << i << " ";
+  log_ << "\n";
+  log_ << "eta_bins: ";
+  for (auto i : this->eta_bins_)
+    log_ << i << " ";
+  log_ << "\n";
+  log_ << "res_bins: ";
+  for (auto i : this->res_bins_)
+    log_ << i << " ";
+  log_ << "\n";
+  edm::LogPrint("NNVtxAssoc") << log_.str();
+}
+
+template bool NNVtxAssoc::TTTrackNetworkSelector<const l1ct::TkObjEmu>(const PFRegionEmu&,
+                                                                       const l1ct::TkObjEmu&,
+                                                                       const l1ct::PVObjEmu&);
+template bool NNVtxAssoc::TTTrackNetworkSelector<const l1ct::PFChargedObjEmu>(const PFRegionEmu&,
+                                                                              const l1ct::PFChargedObjEmu&,
+                                                                              const l1ct::PVObjEmu&);