Add ParT and UParT SONIC producer second pull request

Configuration/ProcessModifiers/python/allSonicTriton_cff.py

@@ -5,6 +5,9 @@
 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
+from Configuration.ProcessModifiers.unifiedparticleTransformerAK4SonicTriton_cff import unifiedparticleTransformerAK4SonicTriton
+
 
 # collect all SonicTriton-related process modifiers here
-allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton)
+allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton,unifiedparticleTransformerAK4SonicTriton)

Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py

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

Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py

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

RecoBTag/ONNXRuntime/BuildFile.xml

@@ -1,3 +1,4 @@
+<use name="PhysicsTools/ONNXRuntime"/>
 <use name="DataFormats/BTauReco"/>
 <export>
   <lib name="1"/>

RecoBTag/ONNXRuntime/interface/tensor_configs.h

@@ -1,6 +1,8 @@
 #ifndef RecoBTag_ONNXRuntime_tensor_configs_h
 #define RecoBTag_ONNXRuntime_tensor_configs_h
 
+#include <map>
+
 namespace deepflavour {
 
   constexpr unsigned n_features_global = 15;
@@ -28,4 +30,85 @@ namespace deepvertex {
 
 }  // namespace deepvertex
 
+namespace parT {
+
+  enum InputFeatures {
+    kBegin=0,
+    kChargedCandidates=kBegin,
+    kNeutralCandidates=1,
+    kVertices=2,
+    kChargedCandidates4Vec=3,
+    kNeutralCandidates4Vec=4,
+    kVertices4Vec=5,
+    kEnd=6
+  };
+
+  constexpr unsigned n_cpf_accept = 25;
+  constexpr unsigned n_npf_accept = 25;
+  constexpr unsigned n_sv_accept = 5;
+
+  const std::map<InputFeatures, unsigned int> N_InputFeatures{
+    {kChargedCandidates, 16},
+    {kNeutralCandidates, 8},
+    {kVertices, 14},
+    {kChargedCandidates4Vec, 4},
+    {kNeutralCandidates4Vec, 4},
+    {kVertices4Vec, 4}
+  };
+
+  const std::map<InputFeatures, unsigned int> N_AcceptedFeatures{
+    {kChargedCandidates, n_cpf_accept},
+    {kNeutralCandidates, n_npf_accept},
+    {kVertices, n_sv_accept},
+    {kChargedCandidates4Vec, n_cpf_accept},
+    {kNeutralCandidates4Vec, n_npf_accept},
+    {kVertices4Vec, n_sv_accept}
+  };
+
+} // namespace parT
+
+
+namespace UparT {
+
+  enum InputFeatures {
+    kBegin=0,
+    kChargedCandidates=kBegin,
+    kLostTracks = 1,
+    kNeutralCandidates=2,
+    kVertices=3,
+    kChargedCandidates4Vec=4,
+    kLostTracks4Vec=5,
+    kNeutralCandidates4Vec=6,
+    kVertices4Vec=7,
+    kEnd=8
+  };
+
+  constexpr unsigned n_cpf_accept = 29;
+  constexpr unsigned n_lt_accept = 5;
+  constexpr unsigned n_npf_accept = 25;
+  constexpr unsigned n_sv_accept = 5;
+
+  const std::map<InputFeatures, unsigned int> N_InputFeatures{
+    {kChargedCandidates, 25},
+    {kLostTracks, 18},
+    {kNeutralCandidates, 8},
+    {kVertices, 14},
+    {kChargedCandidates4Vec, 4},
+    {kLostTracks4Vec, 4},
+    {kNeutralCandidates4Vec, 4},
+    {kVertices4Vec, 4}
+  };
+
+  const std::map<InputFeatures, unsigned int> N_AcceptedFeatures{
+    {kChargedCandidates, n_cpf_accept},
+    {kLostTracks, n_lt_accept},
+    {kNeutralCandidates, n_npf_accept},
+    {kVertices, n_sv_accept},
+    {kChargedCandidates4Vec, n_cpf_accept},
+    {kLostTracks4Vec, n_lt_accept},
+    {kNeutralCandidates4Vec, n_npf_accept},
+    {kVertices4Vec, n_sv_accept}
+  };
+
+} // namespace UparT
 #endif

RecoBTag/ONNXRuntime/interface/tensor_fillers.h

@@ -2,6 +2,10 @@
 #define RecoBTag_ONNXRuntime_tensor_fillers_h
 
 #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
+#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
+#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h"
+#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
 
 namespace btagbtvdeep {
 
@@ -19,6 +23,79 @@ namespace btagbtvdeep {
 
   void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features);
 
+  std::vector<float> inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, UparT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::LostTracksFeatures& lt_features, UparT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, UparT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, UparT::InputFeatures ifeature);
+
+  template<class parT_features>
+  void parT_tensor_filler(cms::Ort::FloatArrays& data, const parT::InputFeatures ifeature, const std::vector<parT_features>& features, const unsigned int max_n, const float*& start, unsigned offset) {
+    float* ptr = nullptr;
+    for (std::size_t n = 0; n < max_n; n++) {
+      const auto& f = features.at(n);
+      ptr = &data[ifeature][offset + n * parT::N_InputFeatures.at(ifeature)];
+      start = ptr;
+      const std::vector<float>& inputs = inputs_parT(f, ifeature);
+      for (unsigned int i = 0; i < inputs.size(); i++) {
+        *ptr = inputs[i];
+        ++ptr;
+      }
+      if (inputs.size() > 0) --ptr;
+      assert(start + parT::N_InputFeatures.at(ifeature) - 1 == ptr);
+    }
+  }
+
+  template<class parT_features>
+  void parT_tensor_filler(std::vector<float>& vdata, const parT::InputFeatures ifeature, const std::vector<parT_features>& features, const unsigned int target_n) {
+    unsigned int n = std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)parT::N_AcceptedFeatures.at(ifeature));
+    for (unsigned int count = 0; count < n; count++) {
+      const std::vector<float>& inputs = inputs_parT(features.at(count), ifeature);
+      vdata.insert(vdata.end(), inputs.begin(), inputs.end());
+    }
+    unsigned int n_features = parT::N_InputFeatures.at(ifeature);
+    if (n < target_n)
+      vdata.insert(vdata.end(), (target_n - n) * n_features, 0); // Add 0 to unfilled part as padding value
+  }
+
+  template<class UparT_features>
+  void UparT_tensor_filler(cms::Ort::FloatArrays& data, const UparT::InputFeatures ifeature, const std::vector<UparT_features>& features, const unsigned int max_n, const float*& start, unsigned offset) {
+    float* ptr = nullptr;
+    for (std::size_t n = 0; n < max_n; n++) {
+      const auto& f = features.at(n);
+      ptr = &data[ifeature][offset + n * UparT::N_InputFeatures.at(ifeature)];
+      start = ptr;
+      const std::vector<float>& inputs = inputs_UparT(f, ifeature);
+      for (unsigned int i = 0; i < inputs.size(); i++) {
+        *ptr = inputs[i];
+        ++ptr;
+      }
+      if (inputs.size() > 0) --ptr;
+      assert(start + UparT::N_InputFeatures.at(ifeature) - 1 == ptr);
+    }
+  }
+
+  template<class UparT_features>
+  void UparT_tensor_filler(std::vector<float>& vdata, const UparT::InputFeatures ifeature, const std::vector<UparT_features>& features, const unsigned int target_n) {
+    unsigned int n = std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)UparT::N_AcceptedFeatures.at(ifeature));
+    for (unsigned int count = 0; count < n; count++) {
+      const std::vector<float>& inputs = inputs_UparT(features.at(count), ifeature);
+      vdata.insert(vdata.end(), inputs.begin(), inputs.end());
+    }
+    unsigned int n_features = UparT::N_InputFeatures.at(ifeature);
+    if (n < target_n)
+      vdata.insert(vdata.end(), (target_n - n) * n_features, 0); // Add 0 to unfilled part as padding value
+  }
+
+
 }  // namespace btagbtvdeep
 
 #endif