FFTJetProducers/plugins/FFTJetPFPileupCleaner.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    FFTJetProducers
0004 // Class:      FFTJetPFPileupCleaner
0005 //
0006 /**\class FFTJetPFPileupCleaner FFTJetPFPileupCleaner.cc RecoJets/FFTJetProducers/plugins/FFTJetPFPileupCleaner.cc
0007
0008  Description: cleans up a collection of partice flow objects
0009
0010  Implementation:
0011      [Notes on implementation]
0012 */
0013 //
0014 // Original Author:  Igor Volobouev
0015 //         Created:  Thu Jul 14 17:50:33 CDT 2011
0016 //
0017 //
0018 #include <cmath>
0019 #include <climits>
0020 #include <utility>
0021 #include <algorithm>
0022
0023 // framework include files
0024 #include "FWCore/Framework/interface/Frameworkfwd.h"
0025 #include "FWCore/Framework/interface/stream/EDProducer.h"
0026 #include "FWCore/Framework/interface/Event.h"
0027 #include "FWCore/Framework/interface/MakerMacros.h"
0028 #include "FWCore/Utilities/interface/Exception.h"
0029 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0030
0031 #include "DataFormats/Common/interface/View.h"
0032 #include "DataFormats/Common/interface/Handle.h"
0033
0034 #include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"
0035 #include "DataFormats/ParticleFlowCandidate/interface/PFCandidateFwd.h"
0036
0037 #include "DataFormats/VertexReco/interface/Vertex.h"
0038 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0039
0040 #define init_param(type, varname) varname(ps.getParameter<type>(#varname))
0041
0042 //
0043 // class declaration
0044 //
0045 class FFTJetPFPileupCleaner : public edm::stream::EDProducer<> {
0046 public:
0047   explicit FFTJetPFPileupCleaner(const edm::ParameterSet&);
0048   FFTJetPFPileupCleaner() = delete;
0049   FFTJetPFPileupCleaner(const FFTJetPFPileupCleaner&) = delete;
0050   FFTJetPFPileupCleaner& operator=(const FFTJetPFPileupCleaner&) = delete;
0051   ~FFTJetPFPileupCleaner() override;
0052
0053 protected:
0054   // methods
0055   void produce(edm::Event&, const edm::EventSetup&) override;
0056
0057 private:
0058   bool isRemovable(reco::PFCandidate::ParticleType ptype) const;
0059   void setRemovalBit(reco::PFCandidate::ParticleType ptype, bool onOff);
0060   void buildRemovalMask();
0061   bool isAcceptableVtx(reco::VertexCollection::const_iterator iv) const;
0062
0063   reco::VertexRef findSomeVertexWFakes(const edm::Handle<reco::VertexCollection>& vertices,
0064                                        const edm::Handle<reco::VertexCollection>& fakeVertices,
0065                                        const reco::PFCandidate& pfcand,
0066                                        bool* fromFakeSet) const;
0067
0068   edm::InputTag PFCandidates;
0069   edm::InputTag Vertices;
0070   edm::InputTag FakePrimaryVertices;
0071
0072   edm::EDGetTokenT<reco::PFCandidateCollection> PFCandidatesToken;
0073   edm::EDGetTokenT<reco::VertexCollection> VerticesToken;
0074   edm::EDGetTokenT<reco::VertexCollection> FakePrimaryVerticesToken;
0075
0076   // The following, if true, will switch to an algorithm
0077   // which takes a fake primary vertex into account
0078   bool useFakePrimaryVertex;
0079
0080   // The following, if true, will cause association of a candidate
0081   // with some vertex no matter what
0082   bool checkClosestZVertex;
0083
0084   // The following switch will check if the primary vertex
0085   // is a neighbor (in Z) of a track and will keep the
0086   // track if this is so, even if it is not directly associated
0087   // with the primary vertex. This switch is meaningful only if
0088   // "checkClosestZVertex" is true.
0089   bool keepIfPVneighbor;
0090
0091   // The following, if true, will cause removal of candidates
0092   // associated with the main vertex
0093   bool removeMainVertex;
0094
0095   // The following will tell us to remove candidates not associated
0096   // with any vertex
0097   bool removeUnassociated;
0098
0099   // Do we want to reverse the decision?
0100   bool reverseRemovalDecision;
0101
0102   // Flags for removing things. See PFCandidate header
0103   // for particle types.
0104   bool remove_X;          // undefined
0105   bool remove_h;          // charged hadron
0106   bool remove_e;          // electron
0107   bool remove_mu;         // muon
0108   bool remove_gamma;      // photon
0109   bool remove_h0;         // neutral hadron
0110   bool remove_h_HF;       // HF tower identified as a hadron
0111   bool remove_egamma_HF;  // HF tower identified as an EM particle
0112
0113   // Mask for removing things
0114   unsigned removalMask;
0115
0116   // Min and max eta for keeping things
0117   double etaMin;
0118   double etaMax;
0119
0120   // Cut for the vertex Ndof
0121   double vertexNdofCut;
0122
0123   // Cut for the vertex Z
0124   double vertexZmaxCut;
0125
0126   // Cut for the vertex rho
0127   double vertexRhoCut;
0128
0129   // Vector for associating tracks with Z positions of the vertices
0130   mutable std::vector<std::pair<double, unsigned> > zAssoc;
0131 };
0132
0133 //
0134 // constructors and destructor
0135 //
0136 FFTJetPFPileupCleaner::FFTJetPFPileupCleaner(const edm::ParameterSet& ps)
0137     : init_param(edm::InputTag, PFCandidates),
0138       init_param(edm::InputTag, Vertices),
0139       init_param(edm::InputTag, FakePrimaryVertices),
0140       init_param(bool, useFakePrimaryVertex),
0141       init_param(bool, checkClosestZVertex),
0142       init_param(bool, keepIfPVneighbor),
0143       init_param(bool, removeMainVertex),
0144       init_param(bool, removeUnassociated),
0145       init_param(bool, reverseRemovalDecision),
0146       init_param(bool, remove_X),
0147       init_param(bool, remove_h),
0148       init_param(bool, remove_e),
0149       init_param(bool, remove_mu),
0150       init_param(bool, remove_gamma),
0151       init_param(bool, remove_h0),
0152       init_param(bool, remove_h_HF),
0153       init_param(bool, remove_egamma_HF),
0154       removalMask(0),
0155       init_param(double, etaMin),
0156       init_param(double, etaMax),
0157       init_param(double, vertexNdofCut),
0158       init_param(double, vertexZmaxCut),
0159       init_param(double, vertexRhoCut) {
0160   buildRemovalMask();
0161
0162   PFCandidatesToken = consumes<reco::PFCandidateCollection>(PFCandidates);
0163   VerticesToken = consumes<reco::VertexCollection>(Vertices);
0164   if (useFakePrimaryVertex)
0165     FakePrimaryVerticesToken = consumes<reco::VertexCollection>(FakePrimaryVertices);
0166
0167   produces<reco::PFCandidateCollection>();
0168 }
0169
0170 FFTJetPFPileupCleaner::~FFTJetPFPileupCleaner() {}
0171
0172 bool FFTJetPFPileupCleaner::isAcceptableVtx(reco::VertexCollection::const_iterator iv) const {
0173   return !iv->isFake() && static_cast<double>(iv->ndof()) > vertexNdofCut && std::abs(iv->z()) < vertexZmaxCut &&
0174          iv->position().rho() < vertexRhoCut;
0175 }
0176
0177 // ------------ method called to produce the data  ------------
0178 void FFTJetPFPileupCleaner::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0179   // get PFCandidates
0180   auto pOutput = std::make_unique<reco::PFCandidateCollection>();
0181
0182   edm::Handle<reco::PFCandidateCollection> pfCandidates;
0183   iEvent.getByToken(PFCandidatesToken, pfCandidates);
0184
0185   // get vertices
0186   edm::Handle<reco::VertexCollection> vertices;
0187   iEvent.getByToken(VerticesToken, vertices);
0188
0189   edm::Handle<reco::VertexCollection> fakeVertices;
0190   if (useFakePrimaryVertex) {
0191     iEvent.getByToken(FakePrimaryVerticesToken, fakeVertices);
0192     if (!fakeVertices.isValid())
0193       throw cms::Exception("FFTJetBadConfig") << "ERROR in FFTJetPFPileupCleaner:"
0194                                                  " could not find fake vertices"
0195                                               << std::endl;
0196   }
0197
0198   const unsigned ncand = pfCandidates->size();
0199   for (unsigned i = 0; i < ncand; ++i) {
0200     reco::PFCandidatePtr candptr(pfCandidates, i);
0201     bool remove = false;
0202
0203     if (isRemovable(candptr->particleId())) {
0204       bool fromFakeSet = false;
0205       reco::VertexRef vertexref(findSomeVertexWFakes(vertices, fakeVertices, *candptr, &fromFakeSet));
0206       if (vertexref.isNull()) {
0207         // Track is not associated with any vertex
0208         // in any of the vertex sets
0209         remove = removeUnassociated;
0210       } else if (vertexref.key() == 0 && (!useFakePrimaryVertex || fromFakeSet)) {
0211         // Track is associated with the main primary vertex
0212         // However, if we are using fake vertices, this only
0213         // matters if the vertex comes from the fake set. If
0214         // it comes from the real set, remove the track anyway
0215         // because in the combined set the associated vertex
0216         // would not be the main primary vertex.
0217         remove = removeMainVertex;
0218       } else
0219         remove = true;
0220     }
0221
0222     // Check the eta range
0223     if (!remove) {
0224       const double eta = candptr->p4().Eta();
0225       remove = eta < etaMin || eta > etaMax;
0226     }
0227
0228     // Should we remember this candidate?
0229     if (reverseRemovalDecision)
0230       remove = !remove;
0231     if (!remove) {
0232       const reco::PFCandidate& cand = (*pfCandidates)[i];
0233       pOutput->push_back(cand);
0234       pOutput->back().setSourceCandidatePtr(candptr);
0235     }
0236   }
0237
0238   iEvent.put(std::move(pOutput));
0239 }
0240
0241 bool FFTJetPFPileupCleaner::isRemovable(const reco::PFCandidate::ParticleType ptype) const {
0242   const unsigned shift = static_cast<unsigned>(ptype);
0243   assert(shift < 32U);
0244   return removalMask & (1U << shift);
0245 }
0246
0247 void FFTJetPFPileupCleaner::setRemovalBit(const reco::PFCandidate::ParticleType ptype, const bool value) {
0248   const unsigned shift = static_cast<unsigned>(ptype);
0249   assert(shift < 32U);
0250   const unsigned mask = (1U << shift);
0251   if (value)
0252     removalMask |= mask;
0253   else
0254     removalMask &= ~mask;
0255 }
0256
0257 // The following essentially follows the code in PFPileUp.cc,
0258 // with added cut on ndof, vertex Z position, and iteration
0259 // over fakes
0260 reco::VertexRef FFTJetPFPileupCleaner::findSomeVertexWFakes(const edm::Handle<reco::VertexCollection>& vertices,
0261                                                             const edm::Handle<reco::VertexCollection>& fakeVertices,
0262                                                             const reco::PFCandidate& pfcand,
0263                                                             bool* fromFakeSet) const {
0264   typedef reco::VertexCollection::const_iterator IV;
0265   typedef reco::Vertex::trackRef_iterator IT;
0266
0267   *fromFakeSet = false;
0268   reco::TrackBaseRef trackBaseRef(pfcand.trackRef());
0269
0270   size_t iVertex = 0;
0271   unsigned nFoundVertex = 0;
0272   const IV vertend(vertices->end());
0273
0274   {
0275     unsigned index = 0;
0276     double bestweight = 0.0;
0277     for (IV iv = vertices->begin(); iv != vertend; ++iv, ++index)
0278       if (isAcceptableVtx(iv)) {
0279         const reco::Vertex& vtx = *iv;
0280
0281         // loop on tracks in vertices
0282         IT trackend(vtx.tracks_end());
0283         for (IT iTrack = vtx.tracks_begin(); iTrack != trackend; ++iTrack) {
0284           const reco::TrackBaseRef& baseRef = *iTrack;
0285
0286           // one of the tracks in the vertex is the same as
0287           // the track considered in the function
0288           if (baseRef == trackBaseRef) {
0289             // select the vertex for which the track has the highest weight
0290             const double w = vtx.trackWeight(baseRef);
0291             if (w > bestweight) {
0292               bestweight = w;
0293               iVertex = index;
0294               nFoundVertex++;
0295             }
0296           }
0297         }
0298       }
0299   }
0300
0301   if (nFoundVertex > 0) {
0302     if (nFoundVertex != 1)
0303       edm::LogWarning("TrackOnTwoVertex") << "a track is shared by at least two vertices. "
0304                                           << "Used to be an assert";
0305
0306     // Check if we can re-associate this track with one
0307     // of the fake vertices
0308     if (useFakePrimaryVertex) {
0309       const double ztrack = pfcand.vertex().z();
0310       double dzmin = std::abs(ztrack - ((*vertices)[iVertex]).z());
0311
0312       const IV fakeEnd(fakeVertices->end());
0313       unsigned index = 0;
0314       for (IV iv = fakeVertices->begin(); iv != fakeEnd; ++iv, ++index)
0315         if (isAcceptableVtx(iv)) {
0316           const double dz = std::abs(ztrack - iv->z());
0317           if (dz < dzmin) {
0318             dzmin = dz;
0319             iVertex = index;
0320             *fromFakeSet = true;
0321           }
0322         }
0323     }
0324
0325     return reco::VertexRef(*fromFakeSet ? fakeVertices : vertices, iVertex);
0326   }
0327
0328   // optional: as a secondary solution, associate the closest vertex in z
0329   if (checkClosestZVertex) {
0330     const double ztrack = pfcand.vertex().z();
0331     bool foundVertex = false;
0332
0333     if (keepIfPVneighbor) {
0334       // Sort all vertices according to their Z coordinate
0335       zAssoc.clear();
0336       unsigned index = 0;
0337       for (IV iv = vertices->begin(); iv != vertend; ++iv, ++index)
0338         if (isAcceptableVtx(iv))
0339           zAssoc.push_back(std::pair<double, unsigned>(iv->z(), index));
0340       const unsigned numRealVertices = index;
0341
0342       // Mix the fake vertex collection into zAssoc.
0343       // Note that we do not reset "index" before doing this.
0344       if (useFakePrimaryVertex) {
0345         const IV fakeEnd(fakeVertices->end());
0346         for (IV iv = fakeVertices->begin(); iv != fakeEnd; ++iv, ++index)
0347           if (isAcceptableVtx(iv))
0348             zAssoc.push_back(std::pair<double, unsigned>(iv->z(), index));
0349       }
0350
0351       // Check where the track z position fits into this sequence
0352       if (!zAssoc.empty()) {
0353         std::sort(zAssoc.begin(), zAssoc.end());
0354         std::pair<double, unsigned> tPair(ztrack, UINT_MAX);
0355         const unsigned iAbove = std::upper_bound(zAssoc.begin(), zAssoc.end(), tPair) - zAssoc.begin();
0356
0357         // Check whether one of the vertices with indices
0358         // iAbove or (iAbove - 1) is a primary vertex.
0359         // If so, return it. Otherwise return the one
0360         // with closest distance to the track.
0361         unsigned ich[2] = {0U, 0U};
0362         unsigned nch = 1;
0363         if (iAbove) {
0364           ich[0] = iAbove - 1U;
0365           ich[1] = iAbove;
0366           if (iAbove < zAssoc.size())
0367             nch = 2;
0368         }
0369
0370         double dzmin = 1.0e100;
0371         unsigned bestVertexNum = UINT_MAX;
0372         for (unsigned icheck = 0; icheck < nch; ++icheck) {
0373           const unsigned zAssocIndex = ich[icheck];
0374           const unsigned vertexNum = zAssoc[zAssocIndex].second;
0375
0376           if (vertexNum == numRealVertices || (!useFakePrimaryVertex && vertexNum == 0U)) {
0377             bestVertexNum = vertexNum;
0378             break;
0379           }
0380
0381           const double dz = std::abs(ztrack - zAssoc[zAssocIndex].first);
0382           if (dz < dzmin) {
0383             dzmin = dz;
0384             bestVertexNum = vertexNum;
0385           }
0386         }
0387
0388         foundVertex = bestVertexNum < UINT_MAX;
0389         if (foundVertex) {
0390           iVertex = bestVertexNum;
0391           if (iVertex >= numRealVertices) {
0392             *fromFakeSet = true;
0393             iVertex -= numRealVertices;
0394           }
0395         }
0396       }
0397     } else {
0398       // This is a simple association algorithm (from PFPileUp)
0399       // extended to take fake vertices into account
0400       double dzmin = 1.0e100;
0401       unsigned index = 0;
0402       for (IV iv = vertices->begin(); iv != vertend; ++iv, ++index)
0403         if (isAcceptableVtx(iv)) {
0404           const double dz = std::abs(ztrack - iv->z());
0405           if (dz < dzmin) {
0406             dzmin = dz;
0407             iVertex = index;
0408             foundVertex = true;
0409           }
0410         }
0411
0412       if (useFakePrimaryVertex) {
0413         const IV fakeEnd(fakeVertices->end());
0414         index = 0;
0415         for (IV iv = fakeVertices->begin(); iv != fakeEnd; ++iv, ++index)
0416           if (isAcceptableVtx(iv)) {
0417             const double dz = std::abs(ztrack - iv->z());
0418             if (dz < dzmin) {
0419               dzmin = dz;
0420               iVertex = index;
0421               *fromFakeSet = true;
0422               foundVertex = true;
0423             }
0424           }
0425       }
0426     }
0427
0428     if (foundVertex)
0429       return reco::VertexRef(*fromFakeSet ? fakeVertices : vertices, iVertex);
0430   }
0431
0432   return reco::VertexRef();
0433 }
0434
0435 void FFTJetPFPileupCleaner::buildRemovalMask() {
0436   setRemovalBit(reco::PFCandidate::X, remove_X);
0437   setRemovalBit(reco::PFCandidate::h, remove_h);
0438   setRemovalBit(reco::PFCandidate::e, remove_e);
0439   setRemovalBit(reco::PFCandidate::mu, remove_mu);
0440   setRemovalBit(reco::PFCandidate::gamma, remove_gamma);
0441   setRemovalBit(reco::PFCandidate::h0, remove_h0);
0442   setRemovalBit(reco::PFCandidate::h_HF, remove_h_HF);
0443   setRemovalBit(reco::PFCandidate::egamma_HF, remove_egamma_HF);
0444 }
0445
0446 //define this as a plug-in
0447 DEFINE_FWK_MODULE(FFTJetPFPileupCleaner);