PatAlgos/plugins/PATTracksToPackedCandidates.cc

0001 //Based on: PhysicsTools/PatAlgos/plugins/PATLostTracks.cc
0002
0003 // system include files
0004 #include <memory>
0005
0006 // user include files
0007 #include "FWCore/Framework/interface/Frameworkfwd.h"
0008 #include "FWCore/Framework/interface/global/EDProducer.h"
0009
0010 #include "FWCore/Framework/interface/Event.h"
0011 #include "FWCore/Framework/interface/MakerMacros.h"
0012
0013 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0014 #include "FWCore/Utilities/interface/StreamID.h"
0015
0016 #include "DataFormats/TrackReco/interface/Track.h"
0017 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0018 #include "DataFormats/VertexReco/interface/Vertex.h"
0019 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0020 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0021
0022 #include "DataFormats/Candidate/interface/Candidate.h"
0023 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0024 #include "DataFormats/Common/interface/Association.h"
0025
0026 #include <vector>
0027
0028 //
0029 // class declaration
0030 //
0031
0032 class PATTracksToPackedCandidates : public edm::global::EDProducer<> {
0033 public:
0034   explicit PATTracksToPackedCandidates(const edm::ParameterSet&);
0035
0036   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0037
0038   void addPackedCandidate(std::vector<pat::PackedCandidate>& cands,
0039                           const reco::Track trk,
0040                           const reco::VertexRef& pvSlimmed,
0041                           const reco::VertexRefProd& pvSlimmedColl,
0042                           bool passPixelTrackSel) const;
0043
0044 private:
0045   void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
0046
0047   // ----------member data ---------------------------
0048   const edm::EDGetTokenT<reco::TrackCollection> srcTracks_;
0049   const edm::EDGetTokenT<reco::VertexCollection> srcPrimaryVertices_;
0050   const edm::EDGetTokenT<reco::BeamSpot> srcOfflineBeamSpot_;
0051   const double dzSigCut_;
0052   const double dxySigCut_;
0053   const double dzSigHP_;
0054   const double dxySigHP_;
0055   const double ptMax_;
0056   const double ptMin_;
0057   const bool resetHP_;
0058   const int covarianceVersion_;
0059   const int covarianceSchema_;
0060 };
0061
0062 //
0063 // constants, enums and typedefs
0064 //
0065
0066 //
0067 // static data member definitions
0068 //
0069
0070 //
0071 // constructors and destructor
0072 //
0073 PATTracksToPackedCandidates::PATTracksToPackedCandidates(const edm::ParameterSet& iConfig)
0074     : srcTracks_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("srcTracks"))),
0075       srcPrimaryVertices_(consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("srcPrimaryVertices"))),
0076       srcOfflineBeamSpot_(consumes<reco::BeamSpot>(iConfig.getParameter<edm::InputTag>("srcOfflineBeamSpot"))),
0077       dzSigCut_(iConfig.getParameter<double>("dzSigCut")),
0078       dxySigCut_(iConfig.getParameter<double>("dxySigCut")),
0079       dzSigHP_(iConfig.getParameter<double>("dzSigHP")),
0080       dxySigHP_(iConfig.getParameter<double>("dxySigHP")),
0081       ptMax_(iConfig.getParameter<double>("ptMax")),
0082       ptMin_(iConfig.getParameter<double>("ptMin")),
0083       resetHP_(iConfig.getParameter<bool>("resetHP")),
0084       covarianceVersion_(iConfig.getParameter<int>("covarianceVersion")),
0085       covarianceSchema_(iConfig.getParameter<int>("covarianceSchema")) {
0086   produces<std::vector<pat::PackedCandidate>>();
0087   produces<edm::Association<pat::PackedCandidateCollection>>();
0088 }
0089
0090 //
0091 // member functions
0092 //
0093
0094 // ------------ method called to produce the data  ------------
0095 void PATTracksToPackedCandidates::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0096   using namespace edm;
0097   using namespace std;
0098
0099   //track collection
0100   auto tracks = iEvent.getHandle(srcTracks_);
0101
0102   auto outPtrTrksAsCands = std::make_unique<std::vector<pat::PackedCandidate>>();
0103
0104   //vtx collection
0105   auto pvs = iEvent.getHandle(srcPrimaryVertices_);
0106   reco::VertexRef pv(pvs.id());
0107   reco::VertexRefProd pvRefProd(pvs);
0108
0109   //best vertex
0110   double bestvzError;
0111   math::XYZPoint bestvtx;
0112   math::Error<3>::type vtx_cov;
0113   if (!pvs->empty()) {
0114     pv = reco::VertexRef(pvs, 0);
0115     const reco::Vertex& vtx = (*pvs)[0];
0116     bestvzError = vtx.zError();
0117     bestvtx = vtx.position();
0118     vtx_cov = vtx.covariance();
0119   } else {
0120     const auto& bs = iEvent.get(srcOfflineBeamSpot_);
0121     bestvzError = bs.z0Error();
0122     bestvtx = bs.position();
0123     vtx_cov = bs.covariance3D();
0124   }
0125
0126   std::vector<int> mapping(tracks->size(), -1);
0127   int savedCandIndx = 0;
0128   int trkIndx = -1;
0129   for (auto const& trk : *tracks) {
0130     trkIndx++;
0131     double dzvtx = std::abs(trk.dz(bestvtx));
0132     double dxyvtx = std::abs(trk.dxy(bestvtx));
0133     double dzerror = std::hypot(trk.dzError(), bestvzError);
0134     double dxyerror = trk.dxyError(bestvtx, vtx_cov);
0135
0136     if (dzvtx >= dzSigCut_ * dzerror)
0137       continue;
0138     if (dxyvtx >= dxySigCut_ * dxyerror)
0139       continue;
0140     if (trk.pt() >= ptMax_ || trk.pt() <= ptMin_)
0141       continue;
0142
0143     bool passSelection = (dzvtx < dzSigHP_ * dzerror && dxyvtx < dxySigHP_ * dxyerror);
0144
0145     addPackedCandidate(*outPtrTrksAsCands, trk, pv, pvRefProd, passSelection);
0146
0147     //for creating the reco::Track -> pat::PackedCandidate map
0148     mapping[trkIndx] = savedCandIndx;
0149     savedCandIndx++;
0150   }
0151   edm::OrphanHandle<pat::PackedCandidateCollection> oh = iEvent.put(std::move(outPtrTrksAsCands));
0152   auto tk2pc = std::make_unique<edm::Association<pat::PackedCandidateCollection>>(oh);
0153   edm::Association<pat::PackedCandidateCollection>::Filler tk2pcFiller(*tk2pc);
0154   tk2pcFiller.insert(tracks, mapping.begin(), mapping.end());
0155   tk2pcFiller.fill();
0156   iEvent.put(std::move(tk2pc));
0157 }
0158
0159 void PATTracksToPackedCandidates::addPackedCandidate(std::vector<pat::PackedCandidate>& cands,
0160                                                      const reco::Track trk,
0161                                                      const reco::VertexRef& pvSlimmed,
0162                                                      const reco::VertexRefProd& pvSlimmedColl,
0163                                                      bool passPixelTrackSel) const {
0164   const float mass = 0.13957018;
0165
0166   int id = 211 * trk.charge();
0167
0168   reco::Candidate::PolarLorentzVector p4(trk.pt(), trk.eta(), trk.phi(), mass);
0169   cands.emplace_back(p4, trk.vertex(), trk.pt(), trk.eta(), trk.phi(), id, pvSlimmedColl, pvSlimmed.key());
0170
0171   if (resetHP_) {
0172     if (passPixelTrackSel)
0173       cands.back().setTrackHighPurity(true);
0174     else
0175       cands.back().setTrackHighPurity(false);
0176   } else {
0177     if (trk.quality(reco::TrackBase::highPurity))
0178       cands.back().setTrackHighPurity(true);
0179     else
0180       cands.back().setTrackHighPurity(false);
0181   }
0182
0183   cands.back().setTrackProperties(trk, covarianceSchema_, covarianceVersion_);
0184 }
0185
0186 // ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
0187 void PATTracksToPackedCandidates::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0188   //The following says we do not know what parameters are allowed so do no validation
0189   // Please change this to state exactly what you do use, even if it is no parameters
0190   edm::ParameterSetDescription desc;
0191   desc.add<edm::InputTag>("srcTracks", {"hiConformalPixelTracks"});
0192   desc.add<edm::InputTag>("srcPrimaryVertices", {"offlineSlimmedPrimaryVertices"});
0193   desc.add<edm::InputTag>("srcOfflineBeamSpot", {"offlineBeamSpot"})
0194       ->setComment("use BeamSpot if empty vtx collection");
0195   desc.add<double>("dzSigCut", 10.0);
0196   desc.add<double>("dxySigCut", 25.0);
0197   desc.add<double>("dzSigHP", 7.0)->setComment("to set HighPurity flag for pixel tracks");
0198   desc.add<double>("dxySigHP", 20.0)->setComment("to set HighPurity flag for pixel tracks");
0199   desc.add<double>("ptMax", 1.0)->setComment("max pT for pixel tracks - above this will use general tracks");
0200   desc.add<double>("ptMin", 0.3)->setComment("min pT for pixel tracks");
0201   desc.add<bool>("resetHP", true)
0202       ->setComment("pixel tracks do not have HP flag set. Use False if does not want to reset HP flag");
0203   desc.add<int>("covarianceVersion", 0)->setComment("so far: 0 is Phase0, 1 is Phase1");
0204   desc.add<int>("covarianceSchema", 520)->setComment("use less accurate schema - reduce size of collection");
0205   descriptions.addWithDefaultLabel(desc);
0206 }
0207
0208 //define this as a plug-in
0209 DEFINE_FWK_MODULE(PATTracksToPackedCandidates);