Merge pull request #40802 from jshin96/portFrom131XTo126X_pujetid_fix

[PileupJetId, Puppi] Backport of #40762 (Pileup ID input variable fix, puppi weight ValueMap access, optional photon protection for existing puppi weights) to CMSSW_12_6_X

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,10 @@ 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 +514,8 @@ 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>("useBugFix",false);
   desc.add<bool>("clonePackedCands", false);
   desc.add<int>("vtxNdofCut", 4);
   desc.add<double>("vtxZCut", 24);

PhysicsTools/NanoAOD/python/custom_jme_cff.py

@@ -272,7 +272,9 @@ def AddPileUpJetIDVars(proc, jetName="", jetSrc="", jetTableName="", jetTaskName
       vertexes  = "offlineSlimmedPrimaryVertices",
       inputIsCorrected = True,
       applyJec  = False,
-      usePuppi = True if "PUPPI" in jetName.upper() else False
+      usePuppi = True if "PUPPI" in jetName.upper() else False,
+      srcConstituentWeights = "packedPFCandidatespuppi" if "PUPPI" in jetName.upper() else "",
+      useBugFix = True
     )
   )
   getattr(proc,jetTaskName).add(getattr(proc, puJetIdVarsCalculator))

RecoJets/JetProducers/interface/PileupJetIdAlgo.h

@@ -29,7 +29,7 @@ class PileupJetIdAlgo {
   ~PileupJetIdAlgo();
 
   PileupJetIdentifier computeIdVariables(
-      const reco::Jet* jet, float jec, const reco::Vertex*, const reco::VertexCollection&, double rho, bool usePuppi);
+      const reco::Jet* jet, float jec, const reco::Vertex*, const reco::VertexCollection&, double rho, bool usePuppi, edm::ValueMap<float>& constituentWeights, bool applyConstituentWeight, bool useBugFix);
 
   void set(const PileupJetIdentifier&);
   float getMVAval(const std::vector<std::string>&, const std::unique_ptr<const GBRForest>&);

RecoJets/JetProducers/plugins/PileupJetIdProducer.cc

@@ -27,7 +27,9 @@ 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")) {
+      usePuppi_(iConfig.getParameter<bool>("usePuppi")),
+      applyConstituentWeight_(false),
+      useBugFix_(iConfig.getParameter<bool>("useBugFix")) {
   if (residualsFromTxt_) {
     residualsTxt_ = iConfig.getParameter<edm::FileInPath>("residualsTxt");
   }
@@ -40,6 +42,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 +69,10 @@ 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);
+  }
 }
 
 // ------------------------------------------------------------------------------------------
@@ -79,6 +90,11 @@ void PileupJetIdProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSe
   Handle<View<Jet>> jetHandle;
   iEvent.getByToken(input_jet_token_, jetHandle);
   const View<Jet>& jets = *jetHandle;
+  edm::ValueMap<float> constituentWeights;
+  if (!input_constituent_weights_token_.isUninitialized()) {
+    constituentWeights = iEvent.get(input_constituent_weights_token_);
+  }
+
 
   // input variables
   Handle<ValueMap<StoredPileupJetIdentifier>> vmap;
@@ -167,7 +183,8 @@ 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());
+      puIdentifier = ialgo->computeIdVariables(theJet, jec, &(*vtx), *vertexes, rho, gc->usePuppi(), constituentWeights, gc->applyConstituentWeight(), gc->useBugFix());
+
       ids.push_back(puIdentifier);
     } else {
       // Or read it from the value map

RecoJets/JetProducers/plugins/PileupJetIdProducer.h

@@ -62,6 +62,8 @@ class GBRForestsAndConstants {
   bool residualsFromTxt() const { return residualsFromTxt_; }
   edm::FileInPath const& residualsTxt() const { return residualsTxt_; }
   bool usePuppi() const { return usePuppi_; }
+  bool applyConstituentWeight() const { return applyConstituentWeight_; }
+  bool useBugFix() const {return useBugFix_; }
 
 private:
   std::vector<PileupJetIdAlgo::AlgoGBRForestsAndConstants> vAlgoGBRForestsAndConstants_;
@@ -74,6 +76,8 @@ class GBRForestsAndConstants {
   bool residualsFromTxt_;
   edm::FileInPath residualsTxt_;
   bool usePuppi_;
+  bool applyConstituentWeight_;
+  bool useBugFix_;
 };
 
 class PileupJetIdProducer : public edm::stream::EDProducer<edm::GlobalCache<GBRForestsAndConstants>> {
@@ -99,6 +103,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_;

RecoJets/JetProducers/python/PileupJetID_cfi.py

@@ -33,6 +33,8 @@
      inputIsCorrected = cms.bool(False),
      residualsFromTxt = cms.bool(False),
      usePuppi = cms.bool(False),
+     srcConstituentWeights = cms.InputTag(""),
+     useBugFix = cms.bool(False)
 #     residualsTxt     = cms.FileInPath("RecoJets/JetProducers/data/download.url") # must be an existing file
 )