Merge pull request #47173 from AlexDeMoor/UParT-v2

Introduce Unified Particle Transformer v2

PhysicsTools/NanoAOD/python/custom_jme_cff.py

@@ -27,10 +27,11 @@
 from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
 from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import _pfParticleTransformerAK4JetTagsAll
 from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll
+from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll
 bTagDiscriminatorsForAK4 = cms.PSet(foo = cms.vstring(
   bTagDeepJet+
   _pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll+_pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll+
-  _pfParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4JetTagsAll
+  _pfParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4V1JetTagsAll
 ))
 run2_nanoAOD_ANY.toModify(
   bTagDiscriminatorsForAK4,

PhysicsTools/NanoAOD/python/jetsAK4_Puppi_cff.py

@@ -94,6 +94,9 @@
         UParTAK4RegPtRawCorr = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptcorr')>0?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptcorr'):-1",float,precision=10,doc="UnifiedParT universal flavor-aware visible pT regression (no neutrinos), correction relative to raw jet pT"),
         UParTAK4RegPtRawCorrNeutrino = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptnu')>0?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptnu'):-1",float,precision=10,doc="UnifiedParT universal flavor-aware pT regression neutrino correction, relative to visible. Correction relative to raw jet pT"),
         UParTAK4RegPtRawRes = Var("?(bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreshigh')+bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreslow'))>0?0.5*(bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreshigh')-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreslow')):-1",float,precision=10,doc="UnifiedParT universal flavor-aware jet pT resolution estimator, (q84 - q16)/2"),
+        UParTAK4V1RegPtRawCorr = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptcorr')>0?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptcorr'):-1",float,precision=10,doc="UnifiedParT V1 universal flavor-aware visible pT regression (no neutrinos), correction relative to raw jet pT"),
+        UParTAK4V1RegPtRawCorrNeutrino = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptnu')>0?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptnu'):-1",float,precision=10,doc="UnifiedParT V1 universal flavor-aware pT regression neutrino correction, relative to visible. Correction relative to raw jet pT"),
+        UParTAK4V1RegPtRawRes = Var("?(bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreshigh')+bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreslow'))>0?0.5*(bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreshigh')-bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreslow')):-1",float,precision=10,doc="UnifiedParT V1 universal flavor-aware jet pT resolution estimator, (q84 - q16)/2"),
         puIdDisc = Var("userFloat('pileupJetIdPuppi:fullDiscriminant')", float,doc="Pileup ID BDT discriminant with 133X Winter24 PuppiV18 training",precision=10),
         hfsigmaEtaEta = Var("userFloat('hfsigmaEtaEta')",float,doc="sigmaEtaEta for HF jets (noise discriminating variable)",precision=10),
         hfsigmaPhiPhi = Var("userFloat('hfsigmaPhiPhi')",float,doc="sigmaPhiPhi for HF jets (noise discriminating variable)",precision=10),
@@ -151,6 +154,9 @@ def nanoAOD_addDeepInfoAK4(process,addParticleNet,addRobustParTAK4=False,addUnif
         print("Updating process to run UnifiedParTAK4")
         from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll as pfUnifiedParticleTransformerAK4JetTagsAll
         _btagDiscriminators += pfUnifiedParticleTransformerAK4JetTagsAll
+        print("Updating process to run UnifiedParTAK4V1")
+        from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll as pfUnifiedParticleTransformerAK4V1JetTagsAll
+        _btagDiscriminators += pfUnifiedParticleTransformerAK4V1JetTagsAll
     if len(_btagDiscriminators)==0: return process
     print("Will recalculate the following discriminators: "+", ".join(_btagDiscriminators))
     updateJetCollection(
@@ -166,6 +172,7 @@ def nanoAOD_addDeepInfoAK4(process,addParticleNet,addRobustParTAK4=False,addUnif
     return process
 
 nanoAOD_addDeepInfoAK4_switch = cms.PSet(
+
     nanoAOD_addParticleNet_switch = cms.untracked.bool(False),
     nanoAOD_addRobustParTAK4Tag_switch = cms.untracked.bool(False),
     nanoAOD_addUnifiedParTAK4Tag_switch = cms.untracked.bool(False)

PhysicsTools/PatAlgos/python/recoLayer0/bTagging_cff.py

@@ -45,6 +45,8 @@
   , 'pfParticleTransformerAK4TagInfos'
     # UnifiedParticleTransformerAK4 tag infos
   , 'pfUnifiedParticleTransformerAK4TagInfos'
+    # UnifiedParticleTransformerAK4V1 tag infos
+  , 'pfUnifiedParticleTransformerAK4V1TagInfos'
     # GlobalParticleTransformerAK8 tag infos
   , 'pfGlobalParticleTransformerAK8TagInfos'
     # DeepDoubleB/C tag infos
@@ -367,6 +369,17 @@
     supportedMetaDiscr[disc] = _pfUnifiedParticleTransformerAK4JetTagsProbs
 # -----------------------------------
 
+# -----------------------------------
+# setup UnifiedParticleTransformer AK4 V1
+from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsProbs, _pfUnifiedParticleTransformerAK4V1JetTagsMetaDiscrs
+# update supportedBtagDiscr
+for disc in _pfUnifiedParticleTransformerAK4V1JetTagsProbs + _pfUnifiedParticleTransformerAK4V1JetTagsMetaDiscrs:
+    supportedBtagDiscr[disc] = [["pfUnifiedParticleTransformerAK4V1TagInfos"]]
+# update supportedMetaDiscr
+for disc in _pfUnifiedParticleTransformerAK4V1JetTagsMetaDiscrs:
+    supportedMetaDiscr[disc] = _pfUnifiedParticleTransformerAK4V1JetTagsProbs
+# -----------------------------------
+
 # -----------------------------------
 # setup Negative ParticleTransformer AK4
 from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import _pfNegativeParticleTransformerAK4JetTagsProbs

PhysicsTools/PatAlgos/python/slimming/applyDeepBtagging_cff.py

@@ -13,7 +13,8 @@ def applyDeepBtagging(process, postfix=""):
     from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll as pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
     from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4CHSCentralJetTagsAll as pfParticleNetFromMiniAODAK4CHSCentralJetTagsAll
     from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll as pfUnifiedParticleTransformerAK4JetTagsAll
-
+    from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll as pfUnifiedParticleTransformerAK4V1JetTagsAll
+
     # update slimmed jets to include DeepFlavour (keep same name)
     # make clone for DeepFlavour-less slimmed jets, so output name is preserved
     addToProcessAndTask('slimmedJetsNoDeepFlavour', slimmedJets.clone(), process, task)
@@ -56,6 +57,7 @@ def applyDeepBtagging(process, postfix=""):
             + pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll
             + pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
             + pfUnifiedParticleTransformerAK4JetTagsAll
+            + pfUnifiedParticleTransformerAK4V1JetTagsAll
     )
 
     updateJetCollection(

PhysicsTools/PatAlgos/python/tools/jetTools.py

@@ -730,6 +730,37 @@ def setupBTagging(process, jetSource, pfCandidates, explicitJTA, pvSource, svSou
                                       flip = flip),
                                     process, task)
 
+            if 'UnifiedParticleTransformerAK4V1TagInfos' in btagInfo:
+                svUsed = svSource
+                if btagInfo == 'pfNegativeUnifiedParticleTransformerAK4V1TagInfos':
+                    svUsed = cms.InputTag(btagPrefix+'inclusiveCandidateNegativeSecondaryVertices'+labelName+postfix)
+                    flip = True 
+                else:
+                    flip = False
+                # use right input tags when running with RECO PF candidates, which actually
+                # depends of whether jets use "particleFlow"
+                if pfCandidates.value() == 'packedPFCandidates':
+                    puppi_value_map = setupPuppiForPackedPF(process)[0]
+                    vertex_associator = cms.InputTag("")
+                else:
+                    puppi_value_map = cms.InputTag("puppi")
+                    vertex_associator = cms.InputTag("primaryVertexAssociation","original")
+
+                # If this jet is a puppi jet, then set is_weighted_jet to true.
+                is_weighted_jet = False
+                if ('puppi' in jetSource.value().lower()):
+                    is_weighted_jet = True
+                addToProcessAndTask(btagPrefix+btagInfo+labelName+postfix,
+                                    btag.pfUnifiedParticleTransformerAK4V1TagInfos.clone(
+                                        jets = jetSource,
+                                        vertices=pvSource,
+                                        secondary_vertices=svUsed,
+                                        puppi_value_map = puppi_value_map,
+                                        vertex_associator = vertex_associator,
+                                        is_weighted_jet = is_weighted_jet,
+                                        flip = flip),
+                                    process, task)              
+
             if btagInfo == 'pfDeepDoubleXTagInfos':
                 # can only run on PAT jets, so the updater needs to be used
                 if 'updated' not in jetSource.value().lower():

RecoBTag/Configuration/python/RecoBTag_cff.py

@@ -15,6 +15,7 @@
 from RecoBTag.ONNXRuntime.pfParticleNetAK4_cff import *
 from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import *
 from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import *
+from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import *
 from RecoBTag.ONNXRuntime.pfGlobalParticleTransformerAK8_cff import *
 from RecoVertex.AdaptiveVertexFinder.inclusiveVertexing_cff import *
 from RecoBTag.PixelCluster.pixelClusterTagInfos_cfi import *

RecoBTag/FeatureTools/plugins/UnifiedParticleTransformerAK4TagInfoProducer.cc

@@ -159,7 +159,7 @@ void UnifiedParticleTransformerAK4TagInfoProducer::fillDescriptions(edm::Configu
   desc.add<double>("min_candidate_pt", 0.10);
   desc.add<bool>("flip", false);
   desc.add<bool>("sort_cand_by_pt", false);
-  desc.add<bool>("fix_lt_sorting", false);
+  desc.add<bool>("fix_lt_sorting", true);
   desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
   desc.add<edm::InputTag>("losttracks", edm::InputTag("lostTracks"));
   desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));

RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc

@@ -40,6 +40,7 @@ class UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDP
   const edm::EDGetTokenT<TagInfoCollection> src_;
   std::vector<std::string> flav_names_;
   std::vector<std::string> input_names_;
+  bool use_dynamic_axes_ = false;
   std::vector<std::string> output_names_;
 
   enum InputIndexes {
@@ -76,6 +77,8 @@ UnifiedParticleTransformerAK4ONNXJetTagsProducer::UnifiedParticleTransformerAK4O
     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
+      use_dynamic_axes_(iConfig.getParameter<edm::FileInPath>("model_path").fullPath().find("v2.onnx") !=
+                        std::string::npos),
       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
   // get output names from flav_names
   for (const auto& flav_name : flav_names_) {
@@ -89,7 +92,8 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Con
   desc.add<edm::InputTag>("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos"));
   desc.add<std::vector<std::string>>(
       "input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"});
-  desc.add<edm::FileInPath>("model_path", edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V00/UParTAK4.onnx"));
+  desc.add<edm::FileInPath>("model_path",
+                            edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V01/UParTAK4_v2.onnx"));
   desc.add<std::vector<std::string>>("output_names", {"softmax"});
   desc.add<std::vector<std::string>>(
       "flav_names",
@@ -164,11 +168,20 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
     const reco::FeaturesTagInfo<btagbtvdeep::UnifiedParticleTransformerAK4Features> taginfo) {
   const auto& features = taginfo.features();
 
-  /// We require a fixed size due to an ONNX conversion issue (to be improved in the future ?) ///
-  n_cpf_ = (unsigned int)29;
-  n_lt_ = (unsigned int)5;
-  n_npf_ = (unsigned int)25;
-  n_sv_ = (unsigned int)5;
+  if (use_dynamic_axes_) {
+    // Use actual sizes for dynamic axes version
+    n_cpf_ = std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)29);
+    n_lt_ = std::clamp((unsigned int)features.lt_features.size(), (unsigned int)1, (unsigned int)5);
+    n_npf_ = std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)25);
+    n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)5);
+
+  } else {
+    // Use fixed sizes for original version
+    n_cpf_ = (unsigned int)29;
+    n_lt_ = (unsigned int)5;
+    n_npf_ = (unsigned int)25;
+    n_sv_ = (unsigned int)5;
+  }
 
   input_sizes_ = {
       n_cpf_ * n_features_cpf_,