Merge pull request #44641 from AlexDeMoor/UParT

Introduce Unified Particle Transformer AK4 jet tagger

DataFormats/BTauReco/interface/ChargedCandidateFeatures.h

@@ -22,7 +22,6 @@ namespace btagbtvdeep {
     float btagPf_trackSip3dSig;
     float btagPf_trackSip2dVal;
     float btagPf_trackSip2dSig;
-
     float btagPf_trackJetDistVal;
 
     float drsubjet1;
@@ -39,6 +38,15 @@ namespace btagbtvdeep {
     float drminsv;
     float distminsv;
 
+    float charge;
+    float btagPf_trackDecayLen;
+    float HadFrac;
+    float CaloFrac;
+    float pdgID;
+    float lostInnerHits;
+    float numberOfPixelHits;
+    float numberOfStripHits;
+
     float pt;
     float px;
     float py;

DataFormats/BTauReco/interface/LostTracksFeatures.h

@@ -0,0 +1,34 @@
+#ifndef DataFormats_BTauReco_LostTracksFeatures_h
+#define DataFormats_BTauReco_LostTracksFeatures_h
+
+namespace btagbtvdeep {
+
+  class LostTracksFeatures {
+  public:
+    float btagPf_trackEtaRel;
+    float btagPf_trackPtRel;
+    float btagPf_trackPPar;
+    float btagPf_trackDeltaR;
+    float btagPf_trackPParRatio;
+    float btagPf_trackSip2dVal;
+    float btagPf_trackSip2dSig;
+    float btagPf_trackSip3dVal;
+    float btagPf_trackSip3dSig;
+    float btagPf_trackJetDistVal;
+    float drminsv;
+    float charge;
+    float puppiw;
+    float chi2;
+    float quality;
+    float lostInnerHits;
+    float numberOfPixelHits;
+    float numberOfStripHits;
+    float pt;
+    float eta;
+    float phi;
+    float e;
+  };
+
+}  // namespace btagbtvdeep
+
+#endif  //DataFormats_BTauReco_LostTracksFeatures_h

DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h

@@ -0,0 +1,25 @@
+#ifndef DataFormats_BTauReco_UnifiedParticleTransformerAK4Features_h
+#define DataFormats_BTauReco_UnifiedParticleTransformerAK4Features_h
+
+#include "DataFormats/BTauReco/interface/SecondaryVertexFeatures.h"
+#include "DataFormats/BTauReco/interface/NeutralCandidateFeatures.h"
+#include "DataFormats/BTauReco/interface/ChargedCandidateFeatures.h"
+#include "DataFormats/BTauReco/interface/LostTracksFeatures.h"
+
+#include <vector>
+
+namespace btagbtvdeep {
+
+  class UnifiedParticleTransformerAK4Features {
+  public:
+    bool is_filled = true;
+    std::vector<SecondaryVertexFeatures> sv_features;
+
+    std::vector<NeutralCandidateFeatures> n_pf_features;
+    std::vector<ChargedCandidateFeatures> c_pf_features;
+    std::vector<LostTracksFeatures> lt_features;
+  };
+
+}  // namespace btagbtvdeep
+
+#endif  //DataFormats_BTauReco_UnifiedParticleTransformerAK4Features_h

DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4TagInfo.h

@@ -0,0 +1,15 @@
+#ifndef DataFormats_BTauReco_UnifiedParticleTransformerAK4TagInfo_h
+#define DataFormats_BTauReco_UnifiedParticleTransformerAK4TagInfo_h
+
+#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h"
+#include "DataFormats/BTauReco/interface/FeaturesTagInfo.h"
+
+namespace reco {
+
+  typedef FeaturesTagInfo<btagbtvdeep::UnifiedParticleTransformerAK4Features> UnifiedParticleTransformerAK4TagInfo;
+
+  DECLARE_EDM_REFS(UnifiedParticleTransformerAK4TagInfo)
+
+}  // namespace reco
+
+#endif  // DataFormats_BTauReco_UnifiedParticleTransformerAK4TagInfo_h

DataFormats/BTauReco/src/classes.h

@@ -65,6 +65,8 @@
 #include "DataFormats/BTauReco/interface/DeepDoubleXTagInfo.h"
 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4TagInfo.h"
+#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h"
+#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4TagInfo.h"
 #include "DataFormats/BTauReco/interface/DeepBoostedJetTagInfo.h"
 #include "DataFormats/BTauReco/interface/PixelClusterTagInfo.h"
 

DataFormats/BTauReco/src/classes_def.xml

@@ -328,6 +328,8 @@
   <class name="edm::reftobase::RefHolder<reco::DeepFlavourTagInfoRef>" />
   <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::ParticleTransformerAK4TagInfoRef>" />
   <class name="edm::reftobase::RefHolder<reco::ParticleTransformerAK4TagInfoRef>" />
+  <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::UnifiedParticleTransformerAK4TagInfoRef>" />
+  <class name="edm::reftobase::RefHolder<reco::UnifiedParticleTransformerAK4TagInfoRef>" />  
   <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::CombinedTauTagInfoRef>" />
   <class name="edm::reftobase::RefHolder<reco::CombinedTauTagInfoRef>" />
   <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::IsolatedTauTagInfoRef>" />
@@ -365,6 +367,8 @@
   <class name="edm::reftobase::RefHolder<reco::DeepFlavourTagInfoFwdRef>" />
   <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::ParticleTransformerAK4TagInfoFwdRef>" />
   <class name="edm::reftobase::RefHolder<reco::ParticleTransformerAK4TagInfoFwdRef>" />
+  <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::UnifiedParticleTransformerAK4TagInfoFwdRef>" />
+  <class name="edm::reftobase::RefHolder<reco::UnifiedParticleTransformerAK4TagInfoFwdRef>" />
   <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::CombinedTauTagInfoFwdRef>" />
   <class name="edm::reftobase::RefHolder<reco::CombinedTauTagInfoFwdRef>" />
   <class name="edm::reftobase::Holder<reco::BaseTagInfo, reco::IsolatedTauTagInfoFwdRef>" />
@@ -426,7 +430,8 @@
    <version ClassVersion="4" checksum="3657612864"/>
     <version ClassVersion="3" checksum="938442516"/>
   </class>
-  <class name="btagbtvdeep::ChargedCandidateFeatures" ClassVersion="6">
+  <class name="btagbtvdeep::ChargedCandidateFeatures" ClassVersion="7">
+   <version ClassVersion="7" checksum="2760889180"/>
    <version ClassVersion="6" checksum="3294985591"/>
    <version ClassVersion="5" checksum="3385708923"/>
     <version ClassVersion="4" checksum="2173116472"/>
@@ -466,6 +471,14 @@
   <class name="reco::ParticleTransformerAK4TagInfoRefProd"/>
   <class name="reco::ParticleTransformerAK4TagInfoRefVector"/>
   <class name="edm::Wrapper<reco::ParticleTransformerAK4TagInfoCollection>"/>
+
+  <class name="reco::UnifiedParticleTransformerAK4TagInfo"/>
+  <class name="reco::UnifiedParticleTransformerAK4TagInfoCollection"/>
+  <class name="reco::UnifiedParticleTransformerAK4TagInfoRef"/>
+  <class name="reco::UnifiedParticleTransformerAK4TagInfoFwdRef"/>
+  <class name="reco::UnifiedParticleTransformerAK4TagInfoRefProd"/>
+  <class name="reco::UnifiedParticleTransformerAK4TagInfoRefVector"/>
+  <class name="edm::Wrapper<reco::UnifiedParticleTransformerAK4TagInfoCollection>"/>
 
   <class name="reco::DeepDoubleXTagInfo"/>
   <class name="reco::DeepDoubleXTagInfoCollection"/>
@@ -476,6 +489,7 @@
   <class name="edm::Wrapper<reco::DeepDoubleXTagInfoCollection>"/>
 
   <class name="btagbtvdeep::DeepBoostedJetFeatures"/>
+  <class name="btagbtvdeep::LostTracksFeatures"/>
   <class name="reco::DeepBoostedJetTagInfo"/>
   <class name="reco::DeepBoostedJetTagInfoCollection"/>
   <class name="reco::DeepBoostedJetTagInfoRef"/>

PhysicsTools/NanoAOD/python/custom_btv_cff.py

@@ -38,9 +38,8 @@ def update_jets_AK4(process):
         'pfNegativeDeepFlavourJetTags:probc',
         'pfNegativeDeepFlavourJetTags:probuds',
         'pfNegativeDeepFlavourJetTags:probg',
-    ] + pfParticleTransformerAK4JetTagsAll + pfNegativeParticleTransformerAK4JetTagsProbs \
-        + pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll + pfNegativeParticleNetFromMiniAODAK4PuppiCentralJetTagsProbs
-
+    ]        + pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll + pfNegativeParticleNetFromMiniAODAK4PuppiCentralJetTagsProbs
+    # \ #+ pfParticleTransformerAK4JetTagsAll + pfNegativeParticleTransformerAK4JetTagsProbs \
     updateJetCollection(
         process,
         jetSource=cms.InputTag('slimmedJetsPuppi'),
@@ -495,7 +494,7 @@ def add_BTV(process,  addAK4=False, addAK8=False, scheme="btvSF"):
                 get_DeepCSV_vars(),
                 get_DeepJet_outputs(),  # outputs are added in any case, inputs only if requested
                 get_ParticleNetAK4_outputs(),
-                get_ParticleTransformerAK4_outputs(),
+                #get_ParticleTransformerAK4_outputs(),# removed in 2024
             ))
 
         # disable the ParT branches in default jetPuppi table

PhysicsTools/NanoAOD/python/custom_jme_cff.py

@@ -26,10 +26,11 @@
 from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll
 from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
 from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import _pfParticleTransformerAK4JetTagsAll
+from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll
 bTagDiscriminatorsForAK4 = cms.PSet(foo = cms.vstring(
   bTagDeepJet+
   _pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll+_pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll+
-  _pfParticleTransformerAK4JetTagsAll
+  _pfParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4JetTagsAll
 ))
 run2_nanoAOD_ANY.toModify(
   bTagDiscriminatorsForAK4,
@@ -198,6 +199,16 @@
   btagRobustParTAK4CvB = Var("?(pt>=15)&&(bDiscriminator('pfParticleTransformerAK4JetTags:probc')+bDiscriminator('pfParticleTransformerAK4JetTags:probb')+bDiscriminator('pfParticleTransformerAK4JetTags:probbb')+bDiscriminator('pfParticleTransformerAK4JetTags:problepb'))>0?bDiscriminator('pfParticleTransformerAK4JetTags:probc')/(bDiscriminator('pfParticleTransformerAK4JetTags:probc')+bDiscriminator('pfParticleTransformerAK4JetTags:probb')+bDiscriminator('pfParticleTransformerAK4JetTags:probbb')+bDiscriminator('pfParticleTransformerAK4JetTags:problepb')):-1",float,doc="RobustParTAK4 c vs b+bb+lepb discriminator",precision=10),
   btagRobustParTAK4QG  = Var("?(pt>=15)&&(bDiscriminator('pfParticleTransformerAK4JetTags:probg')+bDiscriminator('pfParticleTransformerAK4JetTags:probuds'))>0?bDiscriminator('pfParticleTransformerAK4JetTags:probg')/(bDiscriminator('pfParticleTransformerAK4JetTags:probg')+bDiscriminator('pfParticleTransformerAK4JetTags:probuds')):-1",float,doc="RobustParTAK4 g vs uds discriminator",precision=10),
 )
+UNIFIEDPARTAK4VARS = cms.PSet(
+  btagUParTAK4B = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll'):-1",float,precision=10,doc="UnifiedParTAK4 b vs. udscg"),
+  btagUParTAK4CvL = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL'):-1",float,precision=10,doc="UnifiedParTAK4 c vs. udsg"),
+  btagUParTAK4CvB = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB'):-1",float,precision=10,doc="UnifiedParTAK4 c vs. b"),
+  btagUParTAK4QvG = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'):-1",float,precision=10,doc="UnifiedParTAK4 q (udsbc) vs. g"),
+  btagUParTAK4TauVJet = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:TauVsJet')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:TauVsJet'):-1",float,precision=10,doc="UnifiedParTAK4 tau vs. jet"),
+  UParTAK4RegPtRawCorr = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptcorr')>0?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptcorr'):-1",float,precision=10,doc="UnifiedParTAK4 universal flavor-aware visible pT regression (no neutrinos), correction relative to raw jet pT"),
+  UParTAK4RegPtRawCorrNeutrino = Var("?pt>15 && bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptnu')>0?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptnu'):-1",float,precision=10,doc="UnifiedParTAK4 universal flavor-aware pT regression neutrino correction, relative to visible. To apply full regression, multiply raw jet pT by both UParTAK4RegPtRawCorr and UParTAK4RegPtRawCorrNeutrino."),
+  UParTAK4RegPtRawRes = Var("?pt>15 && 0.5*(bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreshigh')-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreslow')) > 0?0.5*(bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreshigh')-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreslow')):-1",float,precision=10,doc="UnifiedParTAK4 universal flavor-aware jet pT resolution estimator, (q84 - q16)/2"),
+)
 PARTICLENETAK4VARS = cms.PSet(
   particleNetAK4_B = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:BvsAll'):-1",float,doc="ParticleNetAK4 tagger b vs all (udsg, c) discriminator",precision=10),
   particleNetAK4_CvsL = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:CvsL'):-1",float,doc="ParticleNetAK4 tagger c vs udsg discriminator",precision=10),
@@ -436,6 +447,21 @@ def AddRobustParTAK4Scores(proc, jetTableName=""):
 
   return proc
 
+def AddUnifiedParTAK4Scores(proc, jetTableName=""):
+  """
+  Store RobustParTAK4 scores in jetTable
+  """
+
+  getattr(proc, jetTableName).variables.btagUParTAK4B = UNIFIEDPARTAK4VARS.btagUParTAK4B
+  getattr(proc, jetTableName).variables.btagUParTAK4CvL = UNIFIEDPARTAK4VARS.btagUParTAK4CvL
+  getattr(proc, jetTableName).variables.btagUParTAK4CvB = UNIFIEDPARTAK4VARS.btagUParTAK4CvB
+  getattr(proc, jetTableName).variables.btagUParTAK4TauVJet = UNIFIEDPARTAK4VARS.btagUParTAK4TauVJet
+  getattr(proc, jetTableName).variables.UParTAK4RegPtRawCorr = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorr
+  getattr(proc, jetTableName).variables.UParTAK4RegPtRawCorrNeutrino = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorrNeutrino
+  getattr(proc, jetTableName).variables.UParTAK4RegPtRawRes = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawRes
+
+  return proc
+
 def AddParticleNetAK4Scores(proc, jetTableName=""):
   """
   Store ParticleNetAK4 scores in jetTable
@@ -616,6 +642,7 @@ def SavePatJets(proc, jetName, payload, patJetFinalColl, jetTablePrefix, jetTabl
     AddDeepJetGluonLQuarkScores(proc,jetTableName=jetTableName)
     AddParticleNetAK4Scores(proc,jetTableName=jetTableName)
     AddRobustParTAK4Scores(proc,jetTableName=jetTableName)
+    AddUnifiedParTAK4Scores(proc,jetTableName=jetTableName)
 
   return proc
 
@@ -746,7 +773,16 @@ def ReclusterAK4PuppiJets(proc, recoJA, runOnMC):
   proc.jetPuppiTable.variables.btagRobustParTAK4B   = ROBUSTPARTAK4VARS.btagRobustParTAK4B
   proc.jetPuppiTable.variables.btagRobustParTAK4CvL = ROBUSTPARTAK4VARS.btagRobustParTAK4CvL
   proc.jetPuppiTable.variables.btagRobustParTAK4CvB = ROBUSTPARTAK4VARS.btagRobustParTAK4CvB
-
+  #
+  # Save UnifiedParTAK4 b-tagging and c-tagging variables
+  #
+  proc.jetPuppiTable.variables.btagUParTAK4B = UNIFIEDPARTAK4VARS.btagUParTAK4B
+  proc.jetPuppiTable.variables.btagUParTAK4CvL = UNIFIEDPARTAK4VARS.btagUParTAK4CvL
+  proc.jetPuppiTable.variables.btagUParTAK4CvB = UNIFIEDPARTAK4VARS.btagUParTAK4CvB
+  proc.jetPuppiTable.variables.btagUParTAK4TauVJet = UNIFIEDPARTAK4VARS.btagUParTAK4TauVJet
+  proc.jetPuppiTable.variables.UParTAK4RegPtRawCorr = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorr
+  proc.jetPuppiTable.variables.UParTAK4RegPtRawCorrNeutrino = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorrNeutrino
+  proc.jetPuppiTable.variables.UParTAK4RegPtRawRes = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawRes
   #
   # For Run-2 eras, don't need to save the low pt AK4 Puppi jet table for MET
   #
@@ -883,7 +919,7 @@ def ReclusterAK4CHSJets(proc, recoJA, runOnMC):
   # one tagger only.
   #
   for varNames in proc.jetTable.variables.parameterNames_():
-    if "btagDeepFlav" in varNames or "btagRobustParT" in varNames:
+    if "btagDeepFlav" in varNames or "btagRobustParT" in varNames or "btagUParT" in varNames:
       delattr(proc.jetTable.variables, varNames)
 
   proc.jetTable.variables.btagPNetB = Var("?pt>15 && bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:BvsAll')>0?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:BvsAll'):-1",float,precision=10,doc="ParticleNet b vs. udscg")