From f0f6646b1cd5e3ef35bd8aa44e9b78c15af272e1 Mon Sep 17 00:00:00 2001
From: Nurfikri Norjoharuddeen <nurfikri89@gmail.com>
Date: Fri, 20 Mar 2020 19:26:10 +0100
Subject: [PATCH] Add quark-gluon likelihood input variables for AK4PFCHS jets.
 Update jerc variables by removing deprecated vars and add new ones in line
 with the latest jets_cff.py

---
 PhysicsTools/NanoAOD/python/custom_jme_cff.py | 44 ++++++++++++++-----
 .../python/tools/jetCollectionTools.py        |  6 ++-
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/PhysicsTools/NanoAOD/python/custom_jme_cff.py b/PhysicsTools/NanoAOD/python/custom_jme_cff.py
index 53b9df7219fa2..261a3896f05d6 100644
--- a/PhysicsTools/NanoAOD/python/custom_jme_cff.py
+++ b/PhysicsTools/NanoAOD/python/custom_jme_cff.py
@@ -174,8 +174,10 @@
   muEF      = jetTable.variables.muEF,
   rawFactor = jetTable.variables.rawFactor,
   jetId     = jetTable.variables.jetId,
-  jercCHPUF = jetTable.variables.jercCHPUF,
-  jercCHF   = jetTable.variables.jercCHF,
+  chFPV0EF  = jetTable.variables.chFPV0EF,
+  chFPV1EF  = jetTable.variables.chFPV1EF,
+  chFPV2EF  = jetTable.variables.chFPV2EF,
+  chFPV3EF  = jetTable.variables.chFPV3EF,
 )
 
 #============================================
@@ -386,6 +388,20 @@ def AddPileUpJetIDVars(proc):
   proc.jetTable.variables.jetRchg  = Var("userFloat('jetRchg')", float, doc="fraction of jet pT carried by the leading charged constituent", precision= 6) 
   proc.jetTable.variables.nCharged = Var("userInt('nCharged')",  int, doc="number of charged constituents")
 
+def AddQGLTaggerVars(proc):
+  #
+  # Save variables as userFloats and userInts
+  # 
+  proc.updatedJetsWithUserData.userFloats.qgl_axis2 = cms.InputTag("qgtagger:axis2")
+  proc.updatedJetsWithUserData.userFloats.qgl_ptD = cms.InputTag("qgtagger:ptD")
+  proc.updatedJetsWithUserData.userInts.qgl_mult = cms.InputTag("qgtagger:mult")
+  #
+  # Specfiy variables in the jetTable to save in NanoAOD
+  #
+  proc.jetTable.variables.qgl_axis2 = Var("userFloat('qgl_axis2')", float, doc="ellipse minor jet axis (Quark vs Gluon likelihood input variable)", precision= 6) 
+  proc.jetTable.variables.qgl_ptD   = Var("userFloat('qgl_ptD')", float, doc="pT-weighted average pT of constituents (Quark vs Gluon likelihood input variable)", precision= 6) 
+  proc.jetTable.variables.qgl_mult  = Var("userInt('qgl_mult')", int, doc="PF candidates multiplicity (Quark vs Gluon likelihood input variable)") 
+
 def PrepJMECustomNanoAOD(process,runOnMC):
   #
   # Additional variables to AK4GenJets 
@@ -421,15 +437,19 @@ def PrepJMECustomNanoAOD(process,runOnMC):
     maxDR = 0.8,
   )
   process.jetSequence.insert(process.jetSequence.index(process.updatedJetsAK8WithUserData), process.jercVarsFatJet)
-  
-  process.updatedJetsAK8WithUserData.userFloats.jercCHPUF = cms.InputTag(
-    "%s:chargedHadronPUEnergyFraction"  % process.jercVarsFatJet.label()
-  )
-  process.updatedJetsAK8WithUserData.userFloats.jercCHF = cms.InputTag(
-    "%s:chargedHadronCHSEnergyFraction" % process.jercVarsFatJet.label()
+
+  process.updatedJetsAK8WithUserData.userFloats = cms.PSet(
+    chFPV0EF = cms.InputTag("%s:chargedFromPV0EnergyFraction" %process.jercVarsFatJet.label()),
+    chFPV1EF = cms.InputTag("%s:chargedFromPV1EnergyFraction" %process.jercVarsFatJet.label()),
+    chFPV2EF = cms.InputTag("%s:chargedFromPV2EnergyFraction" %process.jercVarsFatJet.label()),
+    chFPV3EF = cms.InputTag("%s:chargedFromPV3EnergyFraction" %process.jercVarsFatJet.label()),
   )
-  process.fatJetTable.variables.jercCHPUF = JETVARS.jercCHPUF
-  process.fatJetTable.variables.jercCHF   = JETVARS.jercCHF
+  
+  process.fatJetTable.variables.chFPV0EF  = JETVARS.chFPV0EF
+  process.fatJetTable.variables.chFPV1EF  = JETVARS.chFPV1EF
+  process.fatJetTable.variables.chFPV2EF  = JETVARS.chFPV2EF
+  process.fatJetTable.variables.chFPV3EF  = JETVARS.chFPV3EF
+
   #
   # 
   # 
@@ -443,6 +463,10 @@ def PrepJMECustomNanoAOD(process,runOnMC):
   # Add variables for pileup jet ID studies.
   #
   AddPileUpJetIDVars(process)
+  #
+  # Add variables for quark guon likelihood tagger studies.
+  #
+  AddQGLTaggerVars(process)
 
   ######################################################################################################################
 
diff --git a/PhysicsTools/PatAlgos/python/tools/jetCollectionTools.py b/PhysicsTools/PatAlgos/python/tools/jetCollectionTools.py
index 0de2aeb5432b6..a6e05e9cc08b8 100644
--- a/PhysicsTools/PatAlgos/python/tools/jetCollectionTools.py
+++ b/PhysicsTools/PatAlgos/python/tools/jetCollectionTools.py
@@ -498,8 +498,10 @@ def addRecoJetCollection(self,
         cms.EDProducer("PATJetUserDataEmbedder",
           src = cms.InputTag(selJet),
           userFloats = cms.PSet(
-            jercCHPUF = cms.InputTag("{}:chargedHadronPUEnergyFraction".format(jercVar)),
-            jercCHF   = cms.InputTag("{}:chargedHadronCHSEnergyFraction".format(jercVar)),
+            chFPV0EF = cms.InputTag("{}:chargedFromPV0EnergyFraction".format(jercVar)),
+            chFPV1EF = cms.InputTag("{}:chargedFromPV1EnergyFraction".format(jercVar)),
+            chFPV2EF = cms.InputTag("{}:chargedFromPV2EnergyFraction".format(jercVar)),
+            chFPV3EF = cms.InputTag("{}:chargedFromPV3EnergyFraction".format(jercVar)),
           ),
           userInts = cms.PSet(
             tightId        = cms.InputTag(tightJetId),