add sonic-ize particle transformer to CMSSW_14_1_0_pre0

Configuration/ProcessModifiers/python/allSonicTriton_cff.py

@@ -5,6 +5,8 @@
 from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton
 from Configuration.ProcessModifiers.deepMETSonicTriton_cff import deepMETSonicTriton
 from Configuration.ProcessModifiers.deepTauSonicTriton_cff import deepTauSonicTriton
+from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton
 
 # collect all SonicTriton-related process modifiers here
-allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton)
+allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton)
+

Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py

@@ -0,0 +1,3 @@
+import FWCore.ParameterSet.Config as cms
+
+particleTransformerAK4SonicTriton = cms.Modifier()

RecoBTag/ONNXRuntime/interface/tensor_configs.h

@@ -28,4 +28,23 @@ namespace deepvertex {
 
 }  // namespace deepvertex
 
+namespace parT {
+
+  enum InputIndexes {
+    kChargedCandidates = 0,
+    kNeutralCandidates = 1,
+    kVertices = 2,
+    kChargedCandidates4Vec = 3,
+    kNeutralCandidates4Vec = 4,
+    kVertices4Vec = 5
+  };
+
+  constexpr unsigned n_features_cpf = 16;
+  constexpr unsigned n_pairwise_features_cpf = 4;
+  constexpr unsigned n_features_npf = 8;
+  constexpr unsigned n_pairwise_features_npf = 4;
+  constexpr unsigned n_features_sv = 14;
+  constexpr unsigned n_pairwise_features_sv = 4;
+} // namespace parT
+
 #endif

RecoBTag/ONNXRuntime/interface/tensor_fillers.h

@@ -2,6 +2,7 @@
 #define RecoBTag_ONNXRuntime_tensor_fillers_h
 
 #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
 
 namespace btagbtvdeep {
 
@@ -18,6 +19,30 @@ namespace btagbtvdeep {
   void seedTrack_tensor_filler(float*& ptr, const btagbtvdeep::SeedingTrackFeatures& seed_features);
 
   void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputIndexes idx);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputIndexes idx);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputIndexes idx);
+
+  template<class parT_features>
+  void parT_tensor_filler(float*& ptr, parT::InputIndexes idx , const parT_features pf) {
+    std::vector<float> inputs;
+    inputs = inputs_parT(pf, idx);
+    for (unsigned int i = 0; i < inputs.size(); i++) {
+      *ptr = inputs[i];
+      ++ptr;
+    }
+    if (inputs.size() > 0) --ptr;
+  }
+
+  template<class parT_features>
+  void parT_tensor_filler(std::vector<float>& vdata, parT::InputIndexes idx , const parT_features pf) {
+    std::vector<float> inputs;
+    inputs = inputs_parT(pf, idx);
+    vdata.insert(vdata.end(), inputs.begin(), inputs.end());
+  }
 
 }  // namespace btagbtvdeep
 

RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc

@@ -16,6 +16,9 @@
 
 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
 
+#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
+
 using namespace cms::Ort;
 
 class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
@@ -27,7 +30,7 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer
 
   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet&);
   static void globalEndJob(const ONNXRuntime*);
-
+ 
 private:
   typedef std::vector<reco::ParticleTransformerAK4TagInfo> TagInfoCollection;
   typedef reco::JetTagCollection JetTagCollection;
@@ -41,24 +44,9 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer
   std::vector<std::string> flav_names_;
   std::vector<std::string> input_names_;
   std::vector<std::string> output_names_;
-
-  enum InputIndexes {
-    kChargedCandidates = 0,
-    kNeutralCandidates = 1,
-    kVertices = 2,
-    kChargedCandidates4Vec = 3,
-    kNeutralCandidates4Vec = 4,
-    kVertices4Vec = 5
-  };
-  unsigned n_cpf_;
-  constexpr static unsigned n_features_cpf_ = 16;
-  constexpr static unsigned n_pairwise_features_cpf_ = 4;
-  unsigned n_npf_;
-  constexpr static unsigned n_features_npf_ = 8;
-  constexpr static unsigned n_pairwise_features_npf_ = 4;
-  unsigned n_sv_;
-  constexpr static unsigned n_features_sv_ = 14;
-  constexpr static unsigned n_pairwise_features_sv_ = 4;
+  unsigned int n_cpf_;
+  unsigned int n_npf_;
+  unsigned int n_sv_;
   std::vector<unsigned> input_sizes_;
   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
 
@@ -84,7 +72,7 @@ void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Configurat
   desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
   desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
   desc.add<edm::FileInPath>("model_path",
-                            edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/RobustParTAK4.onnx"));
+                            edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/modelfile/model.onnx"));
   desc.add<std::vector<std::string>>("output_names", {"softmax"});
   desc.add<std::vector<std::string>>(
       "flav_names", std::vector<std::string>{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
@@ -124,12 +112,12 @@ void ParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEvent, cons
       get_input_sizes(taginfo);
 
       // run prediction with dynamic batch size per event
-      input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)n_features_cpf_},
-                       {(int64_t)1, (int64_t)n_npf_, (int64_t)n_features_npf_},
-                       {(int64_t)1, (int64_t)n_sv_, (int64_t)n_features_sv_},
-                       {(int64_t)1, (int64_t)n_cpf_, (int64_t)n_pairwise_features_cpf_},
-                       {(int64_t)1, (int64_t)n_npf_, (int64_t)n_pairwise_features_npf_},
-                       {(int64_t)1, (int64_t)n_sv_, (int64_t)n_pairwise_features_sv_}};
+      input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::n_features_cpf},
+                       {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::n_features_npf},
+                       {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::n_features_sv},
+                       {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::n_pairwise_features_cpf},
+                       {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::n_pairwise_features_npf},
+                       {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::n_pairwise_features_sv}};
 
       outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];
       assert(outputs.size() == flav_names_.size());
@@ -151,24 +139,21 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
     const reco::FeaturesTagInfo<btagbtvdeep::ParticleTransformerAK4Features> taginfo) {
   const auto& features = taginfo.features();
 
-  unsigned int n_cpf = features.c_pf_features.size();
-  unsigned int n_npf = features.n_pf_features.size();
-  unsigned int n_vtx = features.sv_features.size();
+  n_cpf_ = features.c_pf_features.size();
+  n_npf_ = features.n_pf_features.size();
+  n_sv_ = features.sv_features.size();
 
-  n_cpf_ = std::max((unsigned int)1, n_cpf);
-  n_npf_ = std::max((unsigned int)1, n_npf);
-  n_sv_ = std::max((unsigned int)1, n_vtx);
+  n_cpf_ = std::clamp(n_cpf_, (unsigned int)1, (unsigned int)25);
+  n_npf_ = std::clamp(n_npf_, (unsigned int)1, (unsigned int)25);
+  n_sv_ = std::clamp(n_sv_, (unsigned int)1, (unsigned int)5);
 
-  n_cpf_ = std::min((unsigned int)25, n_cpf_);
-  n_npf_ = std::min((unsigned int)25, n_npf_);
-  n_sv_ = std::min((unsigned int)5, n_sv_);
   input_sizes_ = {
-      n_cpf_ * n_features_cpf_,
-      n_npf_ * n_features_npf_,
-      n_sv_ * n_features_sv_,
-      n_cpf_ * n_pairwise_features_cpf_,
-      n_npf_ * n_pairwise_features_npf_,
-      n_sv_ * n_pairwise_features_sv_,
+      n_cpf_ * parT::n_features_cpf,
+      n_npf_ * parT::n_features_npf,
+      n_sv_ * parT::n_features_sv,
+      n_cpf_ * parT::n_pairwise_features_cpf,
+      n_npf_ * parT::n_pairwise_features_npf,
+      n_sv_ * parT::n_pairwise_features_sv,
   };
   // init data storage
   data_.clear();
@@ -185,110 +170,63 @@ void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::Particl
   unsigned offset = 0;
 
   // c_pf candidates
-  auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
+  const auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
   for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) {
     const auto& c_pf_features = features.c_pf_features.at(c_pf_n);
-    ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_];
+    ptr = &data_[parT::kChargedCandidates][offset + c_pf_n * parT::n_features_cpf];
     start = ptr;
-    *ptr = c_pf_features.btagPf_trackEtaRel;
-    *(++ptr) = c_pf_features.btagPf_trackPtRel;
-    *(++ptr) = c_pf_features.btagPf_trackPPar;
-    *(++ptr) = c_pf_features.btagPf_trackDeltaR;
-    *(++ptr) = c_pf_features.btagPf_trackPParRatio;
-    *(++ptr) = c_pf_features.btagPf_trackSip2dVal;
-    *(++ptr) = c_pf_features.btagPf_trackSip2dSig;
-    *(++ptr) = c_pf_features.btagPf_trackSip3dVal;
-    *(++ptr) = c_pf_features.btagPf_trackSip3dSig;
-    *(++ptr) = c_pf_features.btagPf_trackJetDistVal;
-    *(++ptr) = c_pf_features.ptrel;
-    *(++ptr) = c_pf_features.drminsv;
-    *(++ptr) = c_pf_features.vtx_ass;
-    *(++ptr) = c_pf_features.puppiw;
-    *(++ptr) = c_pf_features.chi2;
-    *(++ptr) = c_pf_features.quality;
-    assert(start + n_features_cpf_ - 1 == ptr);
+    parT_tensor_filler(ptr, parT::kChargedCandidates, c_pf_features);
+    assert(start + parT::n_features_cpf - 1 == ptr);
   }
 
   // n_pf candidates
-  auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
+  const auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
   for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) {
     const auto& n_pf_features = features.n_pf_features.at(n_pf_n);
-    ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_];
+    ptr = &data_[parT::kNeutralCandidates][offset + n_pf_n * parT::n_features_npf];
     start = ptr;
-    *ptr = n_pf_features.ptrel;
-    *(++ptr) = n_pf_features.etarel;
-    *(++ptr) = n_pf_features.phirel;
-    *(++ptr) = n_pf_features.deltaR;
-    *(++ptr) = n_pf_features.isGamma;
-    *(++ptr) = n_pf_features.hadFrac;
-    *(++ptr) = n_pf_features.drminsv;
-    *(++ptr) = n_pf_features.puppiw;
-    assert(start + n_features_npf_ - 1 == ptr);
+    parT_tensor_filler(ptr, parT::kNeutralCandidates, n_pf_features);
+    assert(start + parT::n_features_npf - 1 == ptr);
   }
 
   // sv candidates
-  auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
+  const auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
   for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) {
     const auto& sv_features = features.sv_features.at(sv_n);
-    ptr = &data_[kVertices][offset + sv_n * n_features_sv_];
+    ptr = &data_[parT::kVertices][offset + sv_n * parT::n_features_sv];
     start = ptr;
-    *ptr = sv_features.pt;
-    *(++ptr) = sv_features.deltaR;
-    *(++ptr) = sv_features.mass;
-    *(++ptr) = sv_features.etarel;
-    *(++ptr) = sv_features.phirel;
-    *(++ptr) = sv_features.ntracks;
-    *(++ptr) = sv_features.chi2;
-    *(++ptr) = sv_features.normchi2;
-    *(++ptr) = sv_features.dxy;
-    *(++ptr) = sv_features.dxysig;
-    *(++ptr) = sv_features.d3d;
-    *(++ptr) = sv_features.d3dsig;
-    *(++ptr) = sv_features.costhetasvpv;
-    *(++ptr) = sv_features.enratio;
-    assert(start + n_features_sv_ - 1 == ptr);
+    parT_tensor_filler(ptr, parT::kVertices, sv_features);
+    assert(start + parT::n_features_sv - 1 == ptr);
   }
 
   // cpf pairwise features (4-vectors)
-  auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
+  const auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
   for (std::size_t cpf_n = 0; cpf_n < max_cpf_n; cpf_n++) {
     const auto& cpf_pairwise_features = features.c_pf_features.at(cpf_n);
-    ptr = &data_[kChargedCandidates4Vec][offset + cpf_n * n_pairwise_features_cpf_];
+    ptr = &data_[parT::kChargedCandidates4Vec][offset + cpf_n * parT::n_pairwise_features_cpf];
     start = ptr;
-    *ptr = cpf_pairwise_features.px;
-    *(++ptr) = cpf_pairwise_features.py;
-    *(++ptr) = cpf_pairwise_features.pz;
-    *(++ptr) = cpf_pairwise_features.e;
-
-    assert(start + n_pairwise_features_cpf_ - 1 == ptr);
+    parT_tensor_filler(ptr, parT::kChargedCandidates4Vec, cpf_pairwise_features);
+    assert(start + parT::n_pairwise_features_cpf - 1 == ptr);
   }
 
   // npf pairwise features (4-vectors)
-  auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
+  const auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
   for (std::size_t npf_n = 0; npf_n < max_npf_n; npf_n++) {
     const auto& npf_pairwise_features = features.n_pf_features.at(npf_n);
-    ptr = &data_[kNeutralCandidates4Vec][offset + npf_n * n_pairwise_features_npf_];
+    ptr = &data_[parT::kNeutralCandidates4Vec][offset + npf_n * parT::n_pairwise_features_npf];
     start = ptr;
-    *ptr = npf_pairwise_features.px;
-    *(++ptr) = npf_pairwise_features.py;
-    *(++ptr) = npf_pairwise_features.pz;
-    *(++ptr) = npf_pairwise_features.e;
-
-    assert(start + n_pairwise_features_npf_ - 1 == ptr);
+    parT_tensor_filler(ptr, parT::kNeutralCandidates4Vec, npf_pairwise_features);
+    assert(start + parT::n_pairwise_features_npf - 1 == ptr);
   }
 
   // sv pairwise features (4-vectors)
-  auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_);
+  const auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_);
   for (std::size_t sv_N = 0; sv_N < max_sv_N; sv_N++) {
     const auto& sv_pairwise_features = features.sv_features.at(sv_N);
-    ptr = &data_[kVertices4Vec][offset + sv_N * n_pairwise_features_sv_];
+    ptr = &data_[parT::kVertices4Vec][offset + sv_N * parT::n_pairwise_features_sv];
     start = ptr;
-    *ptr = sv_pairwise_features.px;
-    *(++ptr) = sv_pairwise_features.py;
-    *(++ptr) = sv_pairwise_features.pz;
-    *(++ptr) = sv_pairwise_features.e;
-
-    assert(start + n_pairwise_features_sv_ - 1 == ptr);
+    parT_tensor_filler(ptr, parT::kVertices4Vec, sv_pairwise_features);
+    assert(start + parT::n_pairwise_features_sv - 1 == ptr);
   }
 }