Merge pull request #47702 from y19y19/CMSSW_15_1_X_testUParT

Add UParT_V01 model Sonic producer and fix a bug in miniAOD_tools.py

Author: cmsbuild

Configuration/ProcessModifiers/python/allSonicTriton_cff.py

@@ -6,7 +6,9 @@
 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,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton)
+allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton,unifiedparticleTransformerAK4SonicTriton)
 

Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py

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

PhysicsTools/PatAlgos/python/slimming/miniAOD_tools.py

@@ -583,7 +583,7 @@ def _addUTagToTaus(process, task,
     pp_on_AA.toModify(_uTagToTaus_switches, storePNetCHSjets=False, storeUParTPUPPIjets=False)
     _addUTagToTaus(process, task,
                    storePNetCHSjets = _uTagToTaus_switches.storePNetCHSjets.value(),
-                   storeUParTPUPPIjets = _uTagToTaus_switches.storePNetCHSjets.value(),
+                   storeUParTPUPPIjets = _uTagToTaus_switches.storeUParTPUPPIjets.value(),
                    addGenJet = _uTagToTaus_switches.addGenJet.value()
     )
 

RecoBTag/ONNXRuntime/interface/tensor_configs.h

@@ -66,4 +66,47 @@ namespace parT {
 
 }  // 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
+  };
+
+  inline constexpr unsigned n_cpf_accept = 29;
+  inline constexpr unsigned n_lt_accept = 5;
+  inline constexpr unsigned n_npf_accept = 25;
+  inline constexpr unsigned n_sv_accept = 5;
+
+  constexpr std::array<unsigned int, kEnd> N_InputFeatures{{
+      25,  // kChargedCandidates
+      18,  // kLostTracks
+      8,   // kNeutralCandidates
+      14,  // kVertices
+      4,   // kChargedCandidates4Vec
+      4,   // kLostTracks4Vec
+      4,   // kNeutralCandidates4Vec
+      4,   // kVertices4Vec
+  }};
+
+  constexpr std::array<unsigned int, kEnd> N_AcceptedFeatures{{
+      n_cpf_accept,  // kChargedCandidates
+      n_lt_accept,   // kLostTracks
+      n_npf_accept,  // kNeutralCandidates
+      n_sv_accept,   // kVertices
+      n_cpf_accept,  // kChargedCandidates4Vec
+      n_lt_accept,   // kLostTracks4Vec
+      n_npf_accept,  // kNeutralCandidates4Vec
+      n_sv_accept,   // kVertices4Vec
+  }};
+
+}  // namespace UparT
 #endif

RecoBTag/ONNXRuntime/interface/tensor_fillers.h

@@ -3,6 +3,7 @@
 
 #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"
 
@@ -30,6 +31,17 @@ namespace btagbtvdeep {
 
   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,
@@ -47,7 +59,7 @@ namespace btagbtvdeep {
         *ptr = inputs[i];
         ++ptr;
       }
-      if (inputs.size() > 0)
+      if (!inputs.empty())
         --ptr;
       assert(start + parT::N_InputFeatures.at(ifeature) - 1 == ptr);
     }
@@ -69,6 +81,44 @@ namespace btagbtvdeep {
       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.empty())
+        --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