PatAlgos/plugins/PATRefitVertexProducer.cc

0001 /**
0002   \class PATRefitVertexProducer
0003
0004   This producer is intended to take packedCandidates with tracks associated to
0005   the PV and refit the PV (applying or not) BeamSpot constraint
0006
0007   \autor Michal Bluj, NCBJ Warsaw (and then others)
0008
0009  **/
0010
0011 #include "FWCore/Framework/interface/stream/EDProducer.h"
0012 #include "FWCore/Framework/interface/Event.h"
0013 #include "FWCore/Framework/interface/EventSetup.h"
0014 #include "FWCore/Utilities/interface/InputTag.h"
0015 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0016 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0017 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0018 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0019 #include "DataFormats/TrackReco/interface/Track.h"
0020 #include "DataFormats/VertexReco/interface/Vertex.h"
0021 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0022 #include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
0023 #include "TrackingTools/Records/interface/TransientTrackRecord.h"
0024 #include "RecoVertex/VertexPrimitives/interface/TransientVertex.h"
0025 #include "RecoVertex/AdaptiveVertexFit/interface/AdaptiveVertexFitter.h"
0026
0027 #include <memory>
0028
0029 class PATRefitVertexProducer : public edm::stream::EDProducer<> {
0030 public:
0031   explicit PATRefitVertexProducer(const edm::ParameterSet&);
0032   ~PATRefitVertexProducer() override {}
0033
0034   void produce(edm::Event&, const edm::EventSetup&) override;
0035   static void fillDescriptions(edm::ConfigurationDescriptions&);
0036
0037 private:
0038   //--- utility methods
0039
0040   //--- configuration parameters
0041   edm::EDGetTokenT<std::vector<pat::PackedCandidate> > srcCands_, srcLostTracks_, srcEleKfTracks_;
0042   edm::EDGetTokenT<reco::VertexCollection> srcVertices_;
0043   edm::EDGetTokenT<reco::BeamSpot> srcBeamSpot_;
0044   edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> transTrackBuilderToken_;
0045   bool useBeamSpot_;
0046   bool useLostTracks_;
0047   bool useEleKfTracks_;
0048 };
0049
0050 PATRefitVertexProducer::PATRefitVertexProducer(const edm::ParameterSet& cfg)
0051     : srcCands_(consumes<std::vector<pat::PackedCandidate> >(cfg.getParameter<edm::InputTag>("srcCands"))),
0052       srcLostTracks_(consumes<std::vector<pat::PackedCandidate> >(cfg.getParameter<edm::InputTag>("srcLostTracks"))),
0053       srcEleKfTracks_(consumes<std::vector<pat::PackedCandidate> >(cfg.getParameter<edm::InputTag>("srcEleKfTracks"))),
0054       srcVertices_(consumes<reco::VertexCollection>(cfg.getParameter<edm::InputTag>("srcVertices"))),
0055       srcBeamSpot_(consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("srcBeamSpot"))),
0056       transTrackBuilderToken_(esConsumes(edm::ESInputTag("", "TransientTrackBuilder"))),
0057       useBeamSpot_(cfg.getParameter<bool>("useBeamSpot")),
0058       useLostTracks_(cfg.getParameter<bool>("useLostTracks")),
0059       useEleKfTracks_(cfg.getParameter<bool>("useEleKfTracks")) {
0060   produces<reco::VertexCollection>();
0061 }
0062
0063 void PATRefitVertexProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
0064   // Obtain collections
0065   edm::Handle<std::vector<pat::PackedCandidate> > cands;
0066   evt.getByToken(srcCands_, cands);
0067
0068   edm::Handle<std::vector<pat::PackedCandidate> > lostTrackCands;
0069   if (useLostTracks_)
0070     evt.getByToken(srcLostTracks_, lostTrackCands);
0071
0072   edm::Handle<std::vector<pat::PackedCandidate> > eleKfTrackCands;
0073   if (useEleKfTracks_)
0074     evt.getByToken(srcEleKfTracks_, eleKfTrackCands);
0075
0076   edm::Handle<reco::VertexCollection> vertices;
0077   evt.getByToken(srcVertices_, vertices);
0078   const reco::Vertex& pv = vertices->front();
0079   size_t vtxIdx = 0;
0080
0081   edm::Handle<reco::BeamSpot> beamSpot;
0082   if (useBeamSpot_)
0083     evt.getByToken(srcBeamSpot_, beamSpot);
0084
0085   // Get transient track builder
0086   const TransientTrackBuilder& transTrackBuilder = es.getData(transTrackBuilderToken_);
0087
0088   // Output collection
0089   auto outputVertices = std::make_unique<reco::VertexCollection>();
0090   outputVertices->reserve(1);
0091
0092   // Create a new track collection for vertex refit
0093   std::vector<reco::TransientTrack> transTracks;
0094
0095   // loop over the PFCandidates
0096   for (const auto& cand : (*cands)) {
0097     if (cand.charge() == 0 || cand.vertexRef().isNull())
0098       continue;
0099     if (cand.bestTrack() == nullptr)
0100       continue;
0101     auto key = cand.vertexRef().key();
0102     auto quality = cand.pvAssociationQuality();
0103     if (key != vtxIdx ||
0104         (quality != pat::PackedCandidate::UsedInFitTight && quality != pat::PackedCandidate::UsedInFitLoose))
0105       continue;
0106     if (useEleKfTracks_ && std::abs(cand.pdgId()) == 11)
0107       continue;
0108     transTracks.push_back(transTrackBuilder.build(cand.bestTrack()));
0109   }
0110
0111   // loop over the lostTracks
0112   if (useLostTracks_) {
0113     for (const auto& cand : (*lostTrackCands)) {
0114       if (cand.charge() == 0 || cand.vertexRef().isNull())
0115         continue;
0116       if (cand.bestTrack() == nullptr)
0117         continue;
0118       auto key = cand.vertexRef().key();
0119       auto quality = cand.pvAssociationQuality();
0120       if (key != vtxIdx ||
0121           (quality != pat::PackedCandidate::UsedInFitTight && quality != pat::PackedCandidate::UsedInFitLoose))
0122         continue;
0123       transTracks.push_back(transTrackBuilder.build(cand.bestTrack()));
0124     }
0125   }
0126
0127   // loop over the electronKfTracks
0128   if (useEleKfTracks_) {
0129     for (const auto& cand : (*eleKfTrackCands)) {
0130       if (cand.charge() == 0 || cand.vertexRef().isNull())
0131         continue;
0132       if (cand.bestTrack() == nullptr)
0133         continue;
0134       auto key = cand.vertexRef().key();
0135       auto quality = cand.pvAssociationQuality();
0136       if (key != vtxIdx ||
0137           (quality != pat::PackedCandidate::UsedInFitTight && quality != pat::PackedCandidate::UsedInFitLoose))
0138         continue;
0139       transTracks.push_back(transTrackBuilder.build(cand.bestTrack()));
0140     }
0141   }
0142
0143   // Refit the vertex
0144   TransientVertex transVtx;
0145   reco::Vertex refitPV(pv);  // initialized to the original PV
0146
0147   bool fitOK = true;
0148   if (transTracks.size() >= 3) {
0149     AdaptiveVertexFitter avf;
0150     avf.setWeightThreshold(0.1);  // weight per track, allow almost every fit, else --> exception
0151     if (!useBeamSpot_) {
0152       transVtx = avf.vertex(transTracks);
0153     } else {
0154       transVtx = avf.vertex(transTracks, *beamSpot);
0155     }
0156     if (!transVtx.isValid()) {
0157       fitOK = false;
0158     } else {
0159       //MB: protect against rare cases when transVtx is valid but its postion is ill-defined
0160       if (!std::isfinite(transVtx.position().z()))  //MB: it is enough to check one coordinate (?)
0161         fitOK = false;
0162     }
0163   } else
0164     fitOK = false;
0165   if (fitOK) {
0166     refitPV = transVtx;
0167   }
0168
0169   outputVertices->push_back(refitPV);
0170
0171   evt.put(std::move(outputVertices));
0172 }
0173
0174 void PATRefitVertexProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0175   // patRefitVertexProducer
0176   edm::ParameterSetDescription desc;
0177
0178   desc.add<edm::InputTag>("srcVertices", edm::InputTag("offlineSlimmedPrimaryVertices"));
0179   desc.add<edm::InputTag>("srcCands", edm::InputTag("packedPFCandidates"));
0180   desc.add<edm::InputTag>("srcLostTracks", edm::InputTag("lostTracks"));
0181   desc.add<edm::InputTag>("srcEleKfTracks", edm::InputTag("lostTracks:eleTracks"));
0182   desc.add<edm::InputTag>("srcBeamSpot", edm::InputTag("offlineBeamSpot"));
0183   desc.add<bool>("useBeamSpot", true)->setComment("Refit PV with beam-spot constraint");
0184   desc.add<bool>("useLostTracks", true)
0185       ->setComment("Use collection of tracks not used by PF-candidates, aka lost-tracks");
0186   desc.add<bool>("useEleKfTracks", true)
0187       ->setComment("Use collection of electron KF-tracks instead of GSF-tracks of electron PF-candidates");
0188
0189   descriptions.addWithDefaultLabel(desc);
0190 }
0191
0192 #include "FWCore/Framework/interface/MakerMacros.h"
0193 DEFINE_FWK_MODULE(PATRefitVertexProducer);