Merge pull request #40762 from jshin96/from1300pre4_pujetid_fix

[PileupJetId, Puppi] Pileup ID input variable fix, puppi weight ValueMap access, optional photon protection for existing puppi weights
cms-sw · Feb 21, 2023 · 918a7fc · 918a7fc
2 parents 57ed23b + 283d569
commit 918a7fc
Show file tree

Hide file tree

Showing 7 changed files with 165 additions and 89 deletions.
diff --git a/CommonTools/PileupAlgos/plugins/PuppiProducer.cc b/CommonTools/PileupAlgos/plugins/PuppiProducer.cc
@@ -80,6 +80,7 @@ class PuppiProducer : public edm::stream::EDProducer<> {
   uint fNumOfPUVtxsForCharged;
   double fDZCutForChargedFromPUVtxs;
   bool fUseExistingWeights;
+  bool fApplyPhotonProtectionForExistingWeights;
   bool fClonePackedCands;
   int fVtxNdofCut;
   double fVtxZCut;
@@ -104,6 +105,7 @@ PuppiProducer::PuppiProducer(const edm::ParameterSet& iConfig) {
   fNumOfPUVtxsForCharged = iConfig.getParameter<uint>("NumOfPUVtxsForCharged");
   fDZCutForChargedFromPUVtxs = iConfig.getParameter<double>("DeltaZCutForChargedFromPUVtxs");
   fUseExistingWeights = iConfig.getParameter<bool>("useExistingWeights");
+  fApplyPhotonProtectionForExistingWeights = iConfig.getParameter<bool>("applyPhotonProtectionForExistingWeights");
   fClonePackedCands = iConfig.getParameter<bool>("clonePackedCands");
   fVtxNdofCut = iConfig.getParameter<int>("vtxNdofCut");
   fVtxZCut = iConfig.getParameter<double>("vtxZCut");
@@ -363,9 +365,9 @@ void PuppiProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
           curpupweight = lPack->puppiWeight();
         }
       }
-      // Protect high pT photons (important for gamma to hadronic recoil balance)
-      if ((fPtMaxPhotons > 0) && (lPack->pdgId() == 22) && (std::abs(lPack->eta()) < fEtaMaxPhotons) &&
-          (lPack->pt() > fPtMaxPhotons))
+      // Optional: Protect high pT photons (important for gamma to hadronic recoil balance) for existing weights.
+      if (fApplyPhotonProtectionForExistingWeights && (fPtMaxPhotons > 0) && (lPack->pdgId() == 22) &&
+          (std::abs(lPack->eta()) < fEtaMaxPhotons) && (lPack->pt() > fPtMaxPhotons))
         curpupweight = 1;
       lWeights.push_back(curpupweight);
       lPackCtr++;
@@ -511,6 +513,7 @@ void PuppiProducer::fillDescriptions(edm::ConfigurationDescriptions& description
   desc.add<uint>("NumOfPUVtxsForCharged", 0);
   desc.add<double>("DeltaZCutForChargedFromPUVtxs", 0.2);
   desc.add<bool>("useExistingWeights", false);
+  desc.add<bool>("applyPhotonProtectionForExistingWeights", false);
   desc.add<bool>("clonePackedCands", false);
   desc.add<int>("vtxNdofCut", 4);
   desc.add<double>("vtxZCut", 24);

diff --git a/PhysicsTools/NanoAOD/python/custom_jme_cff.py b/PhysicsTools/NanoAOD/python/custom_jme_cff.py
@@ -274,7 +274,7 @@ def AddPileUpJetIDVars(proc, jetName="", jetSrc="", jetTableName="", jetTaskName
       vertexes  = "offlineSlimmedPrimaryVertices",
       inputIsCorrected = True,
       applyJec  = False,
-      usePuppi = True if "PUPPI" in jetName.upper() else False
+      srcConstituentWeights = "packedPFCandidatespuppi" if "PUPPI" in jetName.upper() else ""
     )
   )
   getattr(proc,jetTaskName).add(getattr(proc, puJetIdVarsCalculator))

diff --git a/RecoJets/JetProducers/interface/PileupJetIdAlgo.h b/RecoJets/JetProducers/interface/PileupJetIdAlgo.h
@@ -28,8 +28,13 @@ class PileupJetIdAlgo {
   PileupJetIdAlgo(AlgoGBRForestsAndConstants const* cache);
   ~PileupJetIdAlgo();
 
-  PileupJetIdentifier computeIdVariables(
-      const reco::Jet* jet, float jec, const reco::Vertex*, const reco::VertexCollection&, double rho, bool usePuppi);
+  PileupJetIdentifier computeIdVariables(const reco::Jet* jet,
+                                         float jec,
+                                         const reco::Vertex*,
+                                         const reco::VertexCollection&,
+                                         double rho,
+                                         edm::ValueMap<float>& constituentWeights,
+                                         bool applyConstituentWeight);
 
   void set(const PileupJetIdentifier&);
   float getMVAval(const std::vector<std::string>&, const std::unique_ptr<const GBRForest>&);

diff --git a/RecoJets/JetProducers/plugins/PileupJetIdProducer.cc b/RecoJets/JetProducers/plugins/PileupJetIdProducer.cc
@@ -27,7 +27,7 @@ GBRForestsAndConstants::GBRForestsAndConstants(edm::ParameterSet const& iConfig)
       applyJec_(iConfig.getParameter<bool>("applyJec")),
       jec_(iConfig.getParameter<std::string>("jec")),
       residualsFromTxt_(iConfig.getParameter<bool>("residualsFromTxt")),
-      usePuppi_(iConfig.getParameter<bool>("usePuppi")) {
+      applyConstituentWeight_(false) {
   if (residualsFromTxt_) {
     residualsTxt_ = iConfig.getParameter<edm::FileInPath>("residualsTxt");
   }
@@ -40,6 +40,11 @@ GBRForestsAndConstants::GBRForestsAndConstants(edm::ParameterSet const& iConfig)
   if (!runMvas_) {
     assert(algos.size() == 1);
   }
+
+  edm::InputTag srcConstituentWeights = iConfig.getParameter<edm::InputTag>("srcConstituentWeights");
+  if (!srcConstituentWeights.label().empty()) {
+    applyConstituentWeight_ = true;
+  }
 }
 
 // ------------------------------------------------------------------------------------------
@@ -62,6 +67,11 @@ PileupJetIdProducer::PileupJetIdProducer(const edm::ParameterSet& iConfig, GBRFo
       consumes<edm::ValueMap<StoredPileupJetIdentifier>>(iConfig.getParameter<edm::InputTag>("jetids"));
   input_rho_token_ = consumes<double>(iConfig.getParameter<edm::InputTag>("rho"));
   parameters_token_ = esConsumes(edm::ESInputTag("", globalCache->jec()));
+
+  edm::InputTag srcConstituentWeights = iConfig.getParameter<edm::InputTag>("srcConstituentWeights");
+  if (!srcConstituentWeights.label().empty()) {
+    input_constituent_weights_token_ = consumes<edm::ValueMap<float>>(srcConstituentWeights);
+  }
 }
 
 // ------------------------------------------------------------------------------------------
@@ -80,6 +90,12 @@ void PileupJetIdProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
   iEvent.getByToken(input_jet_token_, jetHandle);
   const View<Jet>& jets = *jetHandle;
 
+  // Constituent weight (e.g PUPPI) Value Map
+  edm::ValueMap<float> constituentWeights;
+  if (!input_constituent_weights_token_.isUninitialized()) {
+    constituentWeights = iEvent.get(input_constituent_weights_token_);
+  }
+
   // input variables
   Handle<ValueMap<StoredPileupJetIdentifier>> vmap;
   if (!gc->produceJetIds()) {
@@ -167,7 +183,9 @@ void PileupJetIdProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
     PileupJetIdentifier puIdentifier;
     if (gc->produceJetIds()) {
       // Compute the input variables
-      puIdentifier = ialgo->computeIdVariables(theJet, jec, &(*vtx), *vertexes, rho, gc->usePuppi());
+      ////////////////////////////// added PUPPI weight Value Map
+      puIdentifier = ialgo->computeIdVariables(
+          theJet, jec, &(*vtx), *vertexes, rho, constituentWeights, gc->applyConstituentWeight());
       ids.push_back(puIdentifier);
     } else {
       // Or read it from the value map

diff --git a/RecoJets/JetProducers/plugins/PileupJetIdProducer.h b/RecoJets/JetProducers/plugins/PileupJetIdProducer.h
@@ -61,7 +61,7 @@ class GBRForestsAndConstants {
   std::string const& jec() const { return jec_; }
   bool residualsFromTxt() const { return residualsFromTxt_; }
   edm::FileInPath const& residualsTxt() const { return residualsTxt_; }
-  bool usePuppi() const { return usePuppi_; }
+  bool applyConstituentWeight() const { return applyConstituentWeight_; }
 
 private:
   std::vector<PileupJetIdAlgo::AlgoGBRForestsAndConstants> vAlgoGBRForestsAndConstants_;
@@ -73,7 +73,7 @@ class GBRForestsAndConstants {
   std::string jec_;
   bool residualsFromTxt_;
   edm::FileInPath residualsTxt_;
-  bool usePuppi_;
+  bool applyConstituentWeight_;
 };
 
 class PileupJetIdProducer : public edm::stream::EDProducer<edm::GlobalCache<GBRForestsAndConstants>> {
@@ -99,6 +99,8 @@ class PileupJetIdProducer : public edm::stream::EDProducer<edm::GlobalCache<GBRF
   std::unique_ptr<FactorizedJetCorrector> jecCor_;
   std::vector<JetCorrectorParameters> jetCorPars_;
 
+  edm::ValueMap<float> constituentWeights_;
+  edm::EDGetTokenT<edm::ValueMap<float>> input_constituent_weights_token_;
   edm::EDGetTokenT<edm::View<reco::Jet>> input_jet_token_;
   edm::EDGetTokenT<reco::VertexCollection> input_vertex_token_;
   edm::EDGetTokenT<edm::ValueMap<StoredPileupJetIdentifier>> input_vm_pujetid_token_;

diff --git a/RecoJets/JetProducers/python/PileupJetID_cfi.py b/RecoJets/JetProducers/python/PileupJetID_cfi.py
@@ -32,7 +32,7 @@
      applyJec = cms.bool(True),
      inputIsCorrected = cms.bool(False),
      residualsFromTxt = cms.bool(False),
-     usePuppi = cms.bool(False),
+     srcConstituentWeights = cms.InputTag(""),
 #     residualsTxt     = cms.FileInPath("RecoJets/JetProducers/data/download.url") # must be an existing file
 )