CHS and PUPPI track-vertex-association for Run3

CommonTools/ParticleFlow/interface/PFPileUpAlgo.h

@@ -28,6 +28,8 @@ class PFPileUpAlgo {
   inline void setVerbose(bool verbose) { verbose_ = verbose; }
 
   inline void setCheckClosestZVertex(bool val) { checkClosestZVertex_ = val; }
+  inline void setNumOfPUVtxsForCharged(unsigned int val) { fNumOfPUVtxsForCharged_ = val; }
+  inline void setDzCutForChargedFromPUVtxs(double val) { fDzCutForChargedFromPUVtxs_ = val; }
 
   const PFCollection& getPFCandidatesFromPU() const { return pfCandidatesFromPU_; }
 
@@ -38,6 +40,8 @@ class PFPileUpAlgo {
 private:
   /// use the closest z vertex if a track is not in a vertex
   bool checkClosestZVertex_;
+  unsigned int fNumOfPUVtxsForCharged_;
+  double fDzCutForChargedFromPUVtxs_;
 
   /// verbose ?
   bool verbose_;

CommonTools/ParticleFlow/plugins/PFPileUp.cc

@@ -52,6 +52,8 @@ class PFPileUp : public edm::stream::EDProducer<> {
 
   void produce(edm::Event&, const edm::EventSetup&) override;
 
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
 private:
   PFPileUpAlgo pileUpAlgo_;
 
@@ -76,6 +78,8 @@ class PFPileUp : public edm::stream::EDProducer<> {
   edm::EDGetTokenT<edm::ValueMap<int>> tokenVertexAssociationQuality_;
   bool fUseVertexAssociation;
   int vertexAssociationQuality_;
+  unsigned int fNumOfPUVtxsForCharged_;
+  double fDzCutForChargedFromPUVtxs_;
 };
 
 PFPileUp::PFPileUp(const edm::ParameterSet& iConfig) {
@@ -91,20 +95,20 @@ PFPileUp::PFPileUp(const edm::ParameterSet& iConfig) {
     tokenVertexAssociationQuality_ =
         consumes<edm::ValueMap<int>>(iConfig.getParameter<edm::InputTag>("vertexAssociation"));
   }
+  fNumOfPUVtxsForCharged_ = iConfig.getParameter<unsigned int>("NumOfPUVtxsForCharged");
+  fDzCutForChargedFromPUVtxs_ = iConfig.getParameter<double>("DzCutForChargedFromPUVtxs");
 
-  enable_ = iConfig.getParameter<bool>("Enable");
+  enable_ = iConfig.getParameter<bool>("enable");
 
   verbose_ = iConfig.getUntrackedParameter<bool>("verbose", false);
 
-  if (iConfig.exists("checkClosestZVertex")) {
-    checkClosestZVertex_ = iConfig.getParameter<bool>("checkClosestZVertex");
-  } else {
-    checkClosestZVertex_ = false;
-  }
+  checkClosestZVertex_ = iConfig.getParameter<bool>("checkClosestZVertex");
 
   // Configure the algo
   pileUpAlgo_.setVerbose(verbose_);
   pileUpAlgo_.setCheckClosestZVertex(checkClosestZVertex_);
+  pileUpAlgo_.setNumOfPUVtxsForCharged(fNumOfPUVtxsForCharged_);
+  pileUpAlgo_.setDzCutForChargedFromPUVtxs(fDzCutForChargedFromPUVtxs_);
 
   //produces<reco::PFCandidateCollection>();
   produces<PFCollection>();
@@ -192,4 +196,20 @@ void PFPileUp::produce(Event& iEvent, const EventSetup& iSetup) {
   iEvent.put(std::move(pOutput));
   // iEvent.put(std::move(pOutputByValue));
 }
+
+void PFPileUp::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("PFCandidates", edm::InputTag("particleFlowTmpPtrs"));
+  desc.add<edm::InputTag>("Vertices", edm::InputTag("offlinePrimaryVertices"));
+  desc.add<bool>("enable", true);
+  desc.addUntracked<bool>("verbose", false);
+  desc.add<bool>("checkClosestZVertex", true);
+  desc.add<bool>("useVertexAssociation", false);
+  desc.add<int>("vertexAssociationQuality", 0);
+  desc.add<edm::InputTag>("vertexAssociation", edm::InputTag(""));
+  desc.add<unsigned int>("NumOfPUVtxsForCharged", 0);
+  desc.add<double>("DzCutForChargedFromPUVtxs", .2);
+  descriptions.add("pfPileUp", desc);
+}
+
 DEFINE_FWK_MODULE(PFPileUp);

CommonTools/ParticleFlow/python/PFBRECO_cff.py

@@ -7,7 +7,7 @@
 from CommonTools.ParticleFlow.TopProjectors.pfNoPileUp_cfi import *
 pfPileUpIsoPFBRECO = pfPileUp.clone( PFCandidates = 'particleFlowPtrs' )
 from Configuration.ProcessModifiers.pp_on_AA_cff import pp_on_AA
-pp_on_AA.toModify(pfPileUpIsoPFBRECO, Enable = False)
+pp_on_AA.toModify(pfPileUpIsoPFBRECO, enable = False)
 
 pfNoPileUpIsoPFBRECO = pfNoPileUp.clone( topCollection = 'pfPileUpIsoPFBRECO',
                                          bottomCollection = 'particleFlowPtrs')

CommonTools/ParticleFlow/python/Tools/enablePileUpCorrection.py

@@ -7,7 +7,7 @@ def enablePileUpCorrectionInPF2PAT( process, postfix, sequence='PF2PAT'):
 
     # pile up subtraction
     getattr(process,"pfNoPileUp"+postfix).enable = True 
-    getattr(process,"pfPileUp"+postfix).Enable = True 
+    getattr(process,"pfPileUp"+postfix).enable = True 
     getattr(process,"pfPileUp"+postfix).checkClosestZVertex = False 
     getattr(process,"pfPileUp"+postfix).Vertices = 'goodOfflinePrimaryVertices'
 

CommonTools/ParticleFlow/python/pfCHS_cff.py

@@ -1,15 +1,9 @@
 import FWCore.ParameterSet.Config as cms
-from CommonTools.ParticleFlow.pfNoPileUpJME_cff import adapt, pfPileUpJME
-from CommonTools.RecoAlgos.sortedPackedPrimaryVertices_cfi import sortedPackedPrimaryVertices
+from CommonTools.ParticleFlow.pfNoPileUpJME_cff import primaryVertexAssociationJME
 
-packedPrimaryVertexAssociationJME = sortedPackedPrimaryVertices.clone(
-  produceSortedVertices = False,
-  producePileUpCollection  = False,
-  produceNoPileUpCollection = False
-)
-adapt(packedPrimaryVertexAssociationJME)
-
-from CommonTools.ParticleFlow.pfNoPileUpPacked_cfi import pfNoPileUpPacked as _pfNoPileUpPacked
-pfCHS = _pfNoPileUpPacked.clone(
-    vertexAssociationQuality=pfPileUpJME.vertexAssociationQuality
+pfCHS = cms.EDFilter("CandPtrSelector",
+              src = cms.InputTag("packedPFCandidates"),
+              cut = cms.string("fromPV(0)>0 || (vertexRef().key<={} && abs(dz(0))<{})".format(
+                  primaryVertexAssociationJME.assignment.NumOfPUVtxsForCharged.value(),
+                  primaryVertexAssociationJME.assignment.DzCutForChargedFromPUVtxs.value()))
 )

CommonTools/ParticleFlow/python/pfNoPileUpJME_cff.py

@@ -5,35 +5,36 @@
 from CommonTools.ParticleFlow.goodOfflinePrimaryVertices_cfi import goodOfflinePrimaryVertices
 from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation
 
-def adapt(primaryVertexAssociationJME):
-  # options for quality PrimaryDz = 6 (used in PUPPI)
-  primaryVertexAssociationJME.assignment.maxDzSigForPrimaryAssignment = 1e10
-  primaryVertexAssociationJME.assignment.maxDzForPrimaryAssignment = 0.3
-  primaryVertexAssociationJME.assignment.maxDzErrorForPrimaryAssignment = 1e10
-  primaryVertexAssociationJME.assignment.NumOfPUVtxsForCharged = 2
-  primaryVertexAssociationJME.assignment.PtMaxCharged = 20.
-  primaryVertexAssociationJME.assignment.EtaMinUseDz = 2.4
-  primaryVertexAssociationJME.assignment.OnlyUseFirstDz = True
-  from Configuration.Eras.Modifier_phase2_common_cff import phase2_common
-  phase2_common.toModify(
-    primaryVertexAssociationJME.assignment,
-    maxDzForPrimaryAssignment=0.1,
-    EtaMinUseDz = 4.0
-    )
+# The following module implements vertex association for JME.
+# It is not run by default to save computing time (but can be run to investigate alternative vertex identification criteria).
+# Instead its parameters are used as input to faster implementations in pfPileUpJME and pfCHS and puppi.
 primaryVertexAssociationJME = primaryVertexAssociation.clone(vertices = 'goodOfflinePrimaryVertices')
-adapt(primaryVertexAssociationJME)
+primaryVertexAssociationJME.assignment.maxDzSigForPrimaryAssignment = 1e10
+primaryVertexAssociationJME.assignment.maxDzForPrimaryAssignment = 0.3
+primaryVertexAssociationJME.assignment.maxDzErrorForPrimaryAssignment = 1e10
+primaryVertexAssociationJME.assignment.NumOfPUVtxsForCharged = 2
+primaryVertexAssociationJME.assignment.DzCutForChargedFromPUVtxs = 0.2
+primaryVertexAssociationJME.assignment.PtMaxCharged = 20.
+primaryVertexAssociationJME.assignment.EtaMinUseDz = 2.4
+primaryVertexAssociationJME.assignment.OnlyUseFirstDz = True
+from Configuration.Eras.Modifier_phase2_common_cff import phase2_common
+phase2_common.toModify(
+  primaryVertexAssociationJME.assignment,
+  maxDzForPrimaryAssignment=0.1,
+  EtaMinUseDz = 4.0
+  )
 
 pfPileUpJME = _pfPileUp.clone(PFCandidates='particleFlowPtrs',
-                              useVertexAssociation = True,
-                              vertexAssociationQuality = 7,
-                              vertexAssociation = ('primaryVertexAssociationJME','original'),
+                              Vertices = 'goodOfflinePrimaryVertices',
+                              checkClosestZVertex = False,
+                              NumOfPUVtxsForCharged = primaryVertexAssociationJME.assignment.NumOfPUVtxsForCharged,
+                              DzCutForChargedFromPUVtxs = primaryVertexAssociationJME.assignment.DzCutForChargedFromPUVtxs,
                               )
 pfNoPileUpJME = _pfNoPileUp.clone(topCollection = 'pfPileUpJME',
                                   bottomCollection = 'particleFlowPtrs' )
 
 pfNoPileUpJMETask = cms.Task(
     goodOfflinePrimaryVertices,
-    primaryVertexAssociationJME,
     pfPileUpJME,
     pfNoPileUpJME
     )

CommonTools/ParticleFlow/src/PFPileUpAlgo.cc

@@ -22,7 +22,9 @@ void PFPileUpAlgo::process(const PFCollection& pfCandidates, const reco::VertexC
 
     // no associated vertex, or primary vertex
     // not pile-up
-    if (ivertex == -1 || ivertex == 0) {
+    if ((ivertex == -1 || ivertex == 0 ||
+         (fNumOfPUVtxsForCharged_ > 0 && !vertices.empty() && ivertex <= int(fNumOfPUVtxsForCharged_) &&
+          std::abs(cand.vertex().z() - vertices[0].z()) < fDzCutForChargedFromPUVtxs_))) {
       if (verbose_)
         std::cout << "VTX " << i << " " << *(pfCandidates[i]) << std::endl;
       pfCandidatesFromVtx_.push_back(pfCandidates[i]);

CommonTools/PileupAlgos/plugins/PuppiProducer.cc

@@ -72,6 +72,7 @@ class PuppiProducer : public edm::stream::EDProducer<> {
   bool fPuppiNoLep;
   bool fUseFromPVLooseTight;
   bool fUseDZ;
+  bool fUseDZforPileup;
   double fDZCut;
   double fEtaMinUseDZ;
   double fPtMaxCharged;
@@ -95,6 +96,7 @@ PuppiProducer::PuppiProducer(const edm::ParameterSet& iConfig) {
   fPuppiNoLep = iConfig.getParameter<bool>("puppiNoLep");
   fUseFromPVLooseTight = iConfig.getParameter<bool>("UseFromPVLooseTight");
   fUseDZ = iConfig.getParameter<bool>("UseDeltaZCut");
+  fUseDZforPileup = iConfig.getParameter<bool>("UseDeltaZCutForPileup");
   fDZCut = iConfig.getParameter<double>("DeltaZCut");
   fEtaMinUseDZ = iConfig.getParameter<double>("EtaMinUseDeltaZ");
   fPtMaxCharged = iConfig.getParameter<double>("PtMaxCharged");
@@ -281,8 +283,10 @@ void PuppiProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
               pReco.id = 1;
             else if (std::abs(pReco.eta) > fEtaMaxCharged)
               pReco.id = 1;
-            else if ((fUseDZ) && (std::abs(pReco.eta) >= fEtaMinUseDZ))
-              pReco.id = (std::abs(pDZ) < fDZCut) ? 1 : 2;
+            else if ((fUseDZ) && (std::abs(pReco.eta) >= fEtaMinUseDZ) && (std::abs(pDZ) < fDZCut))
+              pReco.id = 1;
+            else if ((fUseDZforPileup) && (std::abs(pReco.eta) >= fEtaMinUseDZ) && (std::abs(pDZ) >= fDZCut))
+              pReco.id = 2;
             else if (fUseFromPVLooseTight && tmpFromPV == 1)
               pReco.id = 2;
             else if (fUseFromPVLooseTight && tmpFromPV == 2)
@@ -322,8 +326,10 @@ void PuppiProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
               pReco.id = 1;
             else if (std::abs(pReco.eta) > fEtaMaxCharged)
               pReco.id = 1;
-            else if ((fUseDZ) && (std::abs(pReco.eta) >= fEtaMinUseDZ))
-              pReco.id = (std::abs(pDZ) < fDZCut) ? 1 : 2;
+            else if ((fUseDZ) && (std::abs(pReco.eta) >= fEtaMinUseDZ) && (std::abs(pDZ) < fDZCut))
+              pReco.id = 1;
+            else if ((fUseDZforPileup) && (std::abs(pReco.eta) >= fEtaMinUseDZ) && (std::abs(pDZ) >= fDZCut))
+              pReco.id = 2;
             else if (fUseFromPVLooseTight && lPack->fromPV() == (pat::PackedCandidate::PVLoose))
               pReco.id = 2;
             else if (fUseFromPVLooseTight && lPack->fromPV() == (pat::PackedCandidate::PVTight))
@@ -489,6 +495,7 @@ void PuppiProducer::fillDescriptions(edm::ConfigurationDescriptions& description
   desc.add<bool>("puppiNoLep", false);
   desc.add<bool>("UseFromPVLooseTight", false);
   desc.add<bool>("UseDeltaZCut", true);
+  desc.add<bool>("UseDeltaZCutForPileup", true);
   desc.add<double>("DeltaZCut", 0.3);
   desc.add<double>("EtaMinUseDeltaZ", 0.);
   desc.add<double>("PtMaxCharged", -1.);

CommonTools/PileupAlgos/python/Puppi_cff.py

@@ -25,15 +25,16 @@
                 )
 
 # from cfipython/, generated by PuppiProducer::fillDescriptions
+from CommonTools.ParticleFlow.pfNoPileUpJME_cff import primaryVertexAssociationJME
 import CommonTools.PileupAlgos.PuppiProducer_cfi as _mod
 puppi = _mod.PuppiProducer.clone(
-                       EtaMinUseDeltaZ = 2.4,
-                       NumOfPUVtxsForCharged = 2,
-                       PtMaxCharged   = 20.,
+                       UseDeltaZCutForPileup = False,
+                       DeltaZCut = primaryVertexAssociationJME.assignment.maxDzForPrimaryAssignment,
+                       EtaMinUseDeltaZ = primaryVertexAssociationJME.assignment.EtaMinUseDz,
+                       PtMaxCharged = primaryVertexAssociationJME.assignment.PtMaxCharged,
+                       NumOfPUVtxsForCharged = primaryVertexAssociationJME.assignment.NumOfPUVtxsForCharged,
+                       DeltaZCutForChargedFromPUVtxs = primaryVertexAssociationJME.assignment.DzCutForChargedFromPUVtxs,
                        PtMaxNeutralsStartSlope = 20.,
-                       useVertexAssociation = True,
-                       vertexAssociationQuality = 6,
-                       vertexAssociation = ('primaryVertexAssociationJME','original'),
                        clonePackedCands   = False, # should only be set to True for MiniAOD
                        algos          = { 
                         0: dict( 
@@ -99,8 +100,6 @@
 
 from Configuration.ProcessModifiers.pp_on_AA_cff import pp_on_AA
 pp_on_AA.toModify(puppi, algos = [])
-from Configuration.ProcessModifiers.run2_miniAOD_pp_on_AA_103X_cff import run2_miniAOD_pp_on_AA_103X
-run2_miniAOD_pp_on_AA_103X.toModify(puppi,useVertexAssociation = False) # because the association is only run on cleanedParticleFlow
 
 puppiNoLep = puppi.clone(
     puppiNoLep = True,