Unify vertex association between CHS and PUPPI

CommonTools/ParticleFlow/plugins/PFNoPileUpPacked.cc

@@ -0,0 +1,91 @@
+// system include files
+#include <memory>
+#include <string>
+
+// user include files
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "DataFormats/Candidate/interface/Candidate.h"
+#include "DataFormats/VertexReco/interface/VertexFwd.h"
+#include "DataFormats/Common/interface/Association.h"
+
+#include "DataFormats/VertexReco/interface/Vertex.h"
+
+#include "FWCore/Utilities/interface/Exception.h"
+#include "FWCore/Framework/interface/EventSetup.h"
+
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+
+using namespace std;
+using namespace edm;
+using namespace reco;
+
+/**\class PFNoPileUpPacked
+\brief Identifies pile-up candidates from a collection of Candidates, and
+produces the corresponding collection of NoPileUpCandidates.
+
+\author Andreas Hinzmann
+\date   May 2021
+
+*/
+
+class PFNoPileUpPacked : public edm::stream::EDProducer<> {
+public:
+  typedef edm::View<reco::Candidate> CandidateView;
+  typedef edm::Association<reco::VertexCollection> CandToVertex;
+
+  explicit PFNoPileUpPacked(const edm::ParameterSet&);
+
+  ~PFNoPileUpPacked() override = default;
+
+  void produce(edm::Event&, const edm::EventSetup&) override;
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+  edm::EDGetTokenT<CandidateView> tokenCandidatesView_;
+  edm::EDGetTokenT<reco::VertexCollection> tokenVertices_;
+  edm::EDGetTokenT<CandToVertex> tokenVertexAssociation_;
+  edm::EDGetTokenT<edm::ValueMap<int>> tokenVertexAssociationQuality_;
+  int vertexAssociationQuality_;
+};
+
+PFNoPileUpPacked::PFNoPileUpPacked(const edm::ParameterSet& iConfig) {
+  tokenCandidatesView_ = consumes<CandidateView>(iConfig.getParameter<InputTag>("candidates"));
+  vertexAssociationQuality_ = iConfig.getParameter<int>("vertexAssociationQuality");
+  tokenVertexAssociation_ = consumes<CandToVertex>(iConfig.getParameter<edm::InputTag>("vertexAssociation"));
+  tokenVertexAssociationQuality_ =
+      consumes<edm::ValueMap<int>>(iConfig.getParameter<edm::InputTag>("vertexAssociation"));
+  produces<edm::PtrVector<reco::Candidate>>();
+}
+
+void PFNoPileUpPacked::produce(Event& iEvent, const EventSetup& iSetup) {
+  unique_ptr<edm::PtrVector<reco::Candidate>> pOutput(new edm::PtrVector<reco::Candidate>);
+  Handle<CandidateView> candidateView;
+  iEvent.getByToken(tokenCandidatesView_, candidateView);
+  const edm::Association<reco::VertexCollection>& associatedPV = iEvent.get(tokenVertexAssociation_);
+  const edm::ValueMap<int>& associationQuality = iEvent.get(tokenVertexAssociationQuality_);
+  for (const auto& p : candidateView->ptrs()) {
+    const reco::VertexRef& PVOrig = associatedPV[p];
+    int quality = associationQuality[p];
+    if (!(PVOrig.isNonnull() && (PVOrig.key() > 0) && (quality >= vertexAssociationQuality_)))
+      pOutput->push_back(p);
+  }
+  iEvent.put(std::move(pOutput));
+}
+
+void PFNoPileUpPacked::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("candidates", edm::InputTag("packedPFCandidates"));
+  desc.add<int>("vertexAssociationQuality", 7);
+  desc.add<edm::InputTag>("vertexAssociation", edm::InputTag("packedPrimaryVertexAssociationJME", "original"));
+  descriptions.addWithDefaultLabel(desc);
+}
+
+DEFINE_FWK_MODULE(PFNoPileUpPacked);

CommonTools/ParticleFlow/plugins/PFPileUp.cc

@@ -1,25 +1,97 @@
-#include "CommonTools/ParticleFlow/plugins/PFPileUp.h"
+// system include files
+#include <memory>
+#include <string>
+
+// user include files
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "DataFormats/ParticleFlowCandidate/interface/PFCandidateFwd.h"
+#include "DataFormats/VertexReco/interface/VertexFwd.h"
+#include "DataFormats/Common/interface/Association.h"
+
+#include "CommonTools/ParticleFlow/interface/PFPileUpAlgo.h"
 
 #include "DataFormats/ParticleFlowCandidate/interface/PileUpPFCandidate.h"
 #include "DataFormats/ParticleFlowCandidate/interface/PileUpPFCandidateFwd.h"
 #include "DataFormats/VertexReco/interface/Vertex.h"
 
 #include "FWCore/Framework/interface/ESHandle.h"
 
-// #include "FWCore/MessageLogger/interface/MessageLogger.h"
 #include "FWCore/Utilities/interface/Exception.h"
 #include "FWCore/Framework/interface/EventSetup.h"
 
 using namespace std;
 using namespace edm;
 using namespace reco;
 
+/**\class PFPileUp
+\brief Identifies pile-up candidates from a collection of PFCandidates, and
+produces the corresponding collection of PileUpCandidates.
+
+\author Colin Bernet
+\date   february 2008
+\updated Florian Beaudette 30/03/2012
+
+*/
+
+class PFPileUp : public edm::stream::EDProducer<> {
+public:
+  typedef std::vector<edm::FwdPtr<reco::PFCandidate>> PFCollection;
+  typedef edm::View<reco::PFCandidate> PFView;
+  typedef std::vector<reco::PFCandidate> PFCollectionByValue;
+  typedef edm::Association<reco::VertexCollection> CandToVertex;
+
+  explicit PFPileUp(const edm::ParameterSet&);
+
+  ~PFPileUp() override;
+
+  void produce(edm::Event&, const edm::EventSetup&) override;
+
+private:
+  PFPileUpAlgo pileUpAlgo_;
+
+  /// PFCandidates to be analyzed
+  edm::EDGetTokenT<PFCollection> tokenPFCandidates_;
+  /// fall-back token
+  edm::EDGetTokenT<PFView> tokenPFCandidatesView_;
+
+  /// vertices
+  edm::EDGetTokenT<reco::VertexCollection> tokenVertices_;
+
+  /// enable PFPileUp selection
+  bool enable_;
+
+  /// verbose ?
+  bool verbose_;
+
+  /// use the closest z vertex if a track is not in a vertex
+  bool checkClosestZVertex_;
+
+  edm::EDGetTokenT<CandToVertex> tokenVertexAssociation_;
+  edm::EDGetTokenT<edm::ValueMap<int>> tokenVertexAssociationQuality_;
+  bool fUseVertexAssociation;
+  int vertexAssociationQuality_;
+};
+
 PFPileUp::PFPileUp(const edm::ParameterSet& iConfig) {
   tokenPFCandidates_ = consumes<PFCollection>(iConfig.getParameter<InputTag>("PFCandidates"));
   tokenPFCandidatesView_ = mayConsume<PFView>(iConfig.getParameter<InputTag>("PFCandidates"));
 
   tokenVertices_ = consumes<VertexCollection>(iConfig.getParameter<InputTag>("Vertices"));
 
+  fUseVertexAssociation = iConfig.getParameter<bool>("useVertexAssociation");
+  vertexAssociationQuality_ = iConfig.getParameter<int>("vertexAssociationQuality");
+  if (fUseVertexAssociation) {
+    tokenVertexAssociation_ = consumes<CandToVertex>(iConfig.getParameter<edm::InputTag>("vertexAssociation"));
+    tokenVertexAssociationQuality_ =
+        consumes<edm::ValueMap<int>>(iConfig.getParameter<edm::InputTag>("vertexAssociation"));
+  }
+
   enable_ = iConfig.getParameter<bool>("Enable");
 
   verbose_ = iConfig.getUntrackedParameter<bool>("verbose", false);
@@ -93,10 +165,22 @@ void PFPileUp::produce(Event& iEvent, const EventSetup& iSetup) {
           "error.");
     }
 
-    pileUpAlgo_.process(*pfCandidatesRef, *vertices);
-    pOutput->insert(
-        pOutput->end(), pileUpAlgo_.getPFCandidatesFromPU().begin(), pileUpAlgo_.getPFCandidatesFromPU().end());
-
+    if (fUseVertexAssociation) {
+      const edm::Association<reco::VertexCollection>& associatedPV = iEvent.get(tokenVertexAssociation_);
+      const edm::ValueMap<int>& associationQuality = iEvent.get(tokenVertexAssociationQuality_);
+      PFCollection pfCandidatesFromPU;
+      for (auto& p : (*pfCandidatesRef)) {
+        const reco::VertexRef& PVOrig = associatedPV[p];
+        int quality = associationQuality[p];
+        if (PVOrig.isNonnull() && (PVOrig.key() > 0) && (quality >= vertexAssociationQuality_))
+          pfCandidatesFromPU.push_back(p);
+      }
+      pOutput->insert(pOutput->end(), pfCandidatesFromPU.begin(), pfCandidatesFromPU.end());
+    } else {
+      pileUpAlgo_.process(*pfCandidatesRef, *vertices);
+      pOutput->insert(
+          pOutput->end(), pileUpAlgo_.getPFCandidatesFromPU().begin(), pileUpAlgo_.getPFCandidatesFromPU().end());
+    }
     // for ( PFCollection::const_iterator byValueBegin = pileUpAlgo_.getPFCandidatesFromPU().begin(),
     // 	    byValueEnd = pileUpAlgo_.getPFCandidatesFromPU().end(), ibyValue = byValueBegin;
     // 	  ibyValue != byValueEnd; ++ibyValue ) {
@@ -108,3 +192,4 @@ void PFPileUp::produce(Event& iEvent, const EventSetup& iSetup) {
   iEvent.put(std::move(pOutput));
   // iEvent.put(std::move(pOutputByValue));
 }
+DEFINE_FWK_MODULE(PFPileUp);

CommonTools/ParticleFlow/plugins/SealModule.cc

@@ -1,15 +1,13 @@
 #include "FWCore/Framework/interface/MakerMacros.h"
 
 #include "CommonTools/ParticleFlow/plugins/PFMET.h"
-#include "CommonTools/ParticleFlow/plugins/PFPileUp.h"
 #include "CommonTools/ParticleFlow/plugins/PFCandidateFwdPtrCollectionFilter.h"
 #include "CommonTools/ParticleFlow/plugins/PFJetFwdPtrProducer.h"
 #include "CommonTools/ParticleFlow/plugins/PFTauFwdPtrProducer.h"
 #include "CommonTools/ParticleFlow/plugins/PFCandidateFromFwdPtrProducer.h"
 #include "CommonTools/ParticleFlow/plugins/DeltaBetaWeights.h"
 
 DEFINE_FWK_MODULE(PFMET);
-DEFINE_FWK_MODULE(PFPileUp);
 
 DEFINE_FWK_MODULE(PFCandidateFwdPtrCollectionStringFilter);
 DEFINE_FWK_MODULE(PFCandidateFwdPtrCollectionPdgIdFilter);

CommonTools/ParticleFlow/python/pfCHS_cff.py

@@ -0,0 +1,15 @@
+import FWCore.ParameterSet.Config as cms
+from CommonTools.ParticleFlow.pfNoPileUpJME_cff import adapt, pfPileUpJME
+from CommonTools.RecoAlgos.sortedPackedPrimaryVertices_cfi import sortedPackedPrimaryVertices
+
+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
+)

CommonTools/ParticleFlow/python/pfNoPileUpJME_cff.py

@@ -2,16 +2,38 @@
 
 from CommonTools.ParticleFlow.pfPileUp_cfi  import pfPileUp as _pfPileUp
 from CommonTools.ParticleFlow.TopProjectors.pfNoPileUp_cfi import pfNoPileUp as _pfNoPileUp
-from CommonTools.ParticleFlow.goodOfflinePrimaryVertices_cfi import *
+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
+    )
+primaryVertexAssociationJME = primaryVertexAssociation.clone(vertices = 'goodOfflinePrimaryVertices')
+adapt(primaryVertexAssociationJME)
 
 pfPileUpJME = _pfPileUp.clone(PFCandidates='particleFlowPtrs',
-                              Vertices = 'goodOfflinePrimaryVertices',
-                              checkClosestZVertex = False )
+                              useVertexAssociation = True,
+                              vertexAssociationQuality = 7,
+                              vertexAssociation = ('primaryVertexAssociationJME','original'),
+                              )
 pfNoPileUpJME = _pfNoPileUp.clone(topCollection = 'pfPileUpJME',
                                   bottomCollection = 'particleFlowPtrs' )
 
 pfNoPileUpJMETask = cms.Task(
     goodOfflinePrimaryVertices,
+    primaryVertexAssociationJME,
     pfPileUpJME,
     pfNoPileUpJME
     )

CommonTools/ParticleFlow/python/pfPileUp_cfi.py

@@ -8,5 +8,8 @@
     # pile-up identification now enabled by default. To be studied for jets
     Enable = cms.bool(True),
     verbose = cms.untracked.bool(False),
-    checkClosestZVertex = cms.bool(True)
+    checkClosestZVertex = cms.bool(True),
+    useVertexAssociation = cms.bool(False),
+    vertexAssociationQuality = cms.int32(0),
+    vertexAssociation = cms.InputTag(''),
     )