RecoUtils/plugins/PF_PU_FirstVertexTracks.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    PF_PU_AssoMap
0004 // Class:      PF_PU_FirstVertexTracks
0005 //
0006 /**\class PF_PU_AssoMap PF_PU_FirstVertexTracks.cc CommonTools/RecoUtils/plugins/PF_PU_FirstVertexTracks.cc
0007
0008   Description: Produces collection of tracks associated to the first vertex based on the pf_pu Association Map
0009 */
0010 //
0011 // Original Author:  Matthias Geisler
0012 //         Created:  Wed Apr 18 14:48:37 CEST 2012
0013 // $Id: PF_PU_FirstVertexTracks.cc,v 1.1 2012/11/21 09:57:30 mgeisler Exp $
0014 //
0015 //
0016 #include "CommonTools/RecoUtils/interface/PF_PU_FirstVertexTracks.h"
0017
0018 // system include files
0019 #include <vector>
0020 #include <string>
0021
0022 // user include files
0023 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0024
0025 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0026 #include "DataFormats/Common/interface/Handle.h"
0027 #include "DataFormats/Common/interface/OneToManyWithQualityGeneric.h"
0028 #include "DataFormats/Common/interface/View.h"
0029
0030 #include "DataFormats/TrackReco/interface/TrackBase.h"
0031 #include "DataFormats/EgammaCandidates/interface/GsfElectron.h"
0032 #include "DataFormats/EgammaCandidates/interface/GsfElectronFwd.h"
0033 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0034
0035 //
0036 // constants, enums and typedefs
0037 //
0038
0039 using namespace edm;
0040 using namespace std;
0041 using namespace reco;
0042
0043 typedef vector<pair<TrackRef, int> > TrackQualityPairVector;
0044
0045 //
0046 // static data member definitions
0047 //
0048
0049 //
0050 // constructors and destructor
0051 //
0052 PF_PU_FirstVertexTracks::PF_PU_FirstVertexTracks(const edm::ParameterSet& iConfig) {
0053   //now do what ever other initialization is needed
0054
0055   input_AssociationType_ = iConfig.getParameter<edm::InputTag>("AssociationType");
0056
0057   token_TrackToVertexAssMap_ = mayConsume<TrackToVertexAssMap>(iConfig.getParameter<InputTag>("AssociationMap"));
0058   token_VertexToTrackAssMap_ = mayConsume<VertexToTrackAssMap>(iConfig.getParameter<InputTag>("AssociationMap"));
0059
0060   token_generalTracksCollection_ = consumes<TrackCollection>(iConfig.getParameter<InputTag>("TrackCollection"));
0061
0062   token_VertexCollection_ = mayConsume<VertexCollection>(iConfig.getParameter<InputTag>("VertexCollection"));
0063
0064   input_MinQuality_ = iConfig.getParameter<int>("MinQuality");
0065
0066   //register your products
0067
0068   if (input_AssociationType_.label() == "TracksToVertex") {
0069     produces<TrackCollection>("T2V");
0070   } else {
0071     if (input_AssociationType_.label() == "VertexToTracks") {
0072       produces<TrackCollection>("V2T");
0073     } else {
0074       if (input_AssociationType_.label() == "Both") {
0075         produces<TrackCollection>("T2V");
0076         produces<TrackCollection>("V2T");
0077       } else {
0078         std::cout << "No correct InputTag for AssociationType!" << std::endl;
0079         std::cout << "Won't produce any TrackCollection!" << std::endl;
0080       }
0081     }
0082   }
0083 }
0084
0085 //
0086 // member functions
0087 //
0088
0089 // ------------ method called to produce the data  ------------
0090 void PF_PU_FirstVertexTracks::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0091   unique_ptr<TrackCollection> t2v_firstvertextracks(new TrackCollection());
0092   unique_ptr<TrackCollection> v2t_firstvertextracks(new TrackCollection());
0093
0094   bool t2vassmap = false;
0095   bool v2tassmap = false;
0096
0097   //get the input vertex<->general track association map
0098   Handle<TrackToVertexAssMap> t2vAM;
0099   Handle<VertexToTrackAssMap> v2tAM;
0100
0101   string asstype = input_AssociationType_.label();
0102
0103   if ((asstype == "TracksToVertex") || (asstype == "Both")) {
0104     if (iEvent.getByToken(token_TrackToVertexAssMap_, t2vAM)) {
0105       t2vassmap = true;
0106     }
0107   }
0108
0109   if ((asstype == "VertexToTracks") || (asstype == "Both")) {
0110     if (iEvent.getByToken(token_VertexToTrackAssMap_, v2tAM)) {
0111       v2tassmap = true;
0112     }
0113   }
0114
0115   if (!t2vassmap && !v2tassmap) {
0116     cout << "No input collection could be found" << endl;
0117     return;
0118   }
0119
0120   //get the input track collection
0121   Handle<TrackCollection> input_trckcollH;
0122   iEvent.getByToken(token_generalTracksCollection_, input_trckcollH);
0123
0124   if (t2vassmap) {
0125     const TrackQualityPairVector trckcoll = t2vAM->begin()->val;
0126
0127     //get the tracks associated to the first vertex and store them in a track collection
0128     for (unsigned int trckcoll_ite = 0; trckcoll_ite < trckcoll.size(); trckcoll_ite++) {
0129       float quality = trckcoll[trckcoll_ite].second;
0130
0131       if (quality >= input_MinQuality_) {
0132         TrackRef AMtrkref = trckcoll[trckcoll_ite].first;
0133
0134         for (unsigned int index_input_trck = 0; index_input_trck < input_trckcollH->size(); index_input_trck++) {
0135           TrackRef input_trackref = TrackRef(input_trckcollH, index_input_trck);
0136
0137           if (trackMatch(*AMtrkref, *input_trackref)) {
0138             t2v_firstvertextracks->push_back(*AMtrkref);
0139             break;
0140           }
0141         }
0142       }
0143     }
0144
0145     iEvent.put(std::move(t2v_firstvertextracks), "T2V");
0146   }
0147
0148   if (v2tassmap) {
0149     //get the input vertex collection
0150     Handle<VertexCollection> input_vtxcollH;
0151     iEvent.getByToken(token_VertexCollection_, input_vtxcollH);
0152
0153     VertexRef firstVertexRef(input_vtxcollH, 0);
0154
0155     VertexToTrackAssMap::const_iterator v2t_ite;
0156
0157     for (v2t_ite = v2tAM->begin(); v2t_ite != v2tAM->end(); v2t_ite++) {
0158       TrackRef AMtrkref = v2t_ite->key;
0159
0160       for (unsigned int index_input_trck = 0; index_input_trck < input_trckcollH->size(); index_input_trck++) {
0161         TrackRef input_trackref = TrackRef(input_trckcollH, index_input_trck);
0162
0163         if (trackMatch(*AMtrkref, *input_trackref)) {
0164           for (unsigned v_ite = 0; v_ite < (v2t_ite->val).size(); v_ite++) {
0165             VertexRef vtxref = (v2t_ite->val)[v_ite].first;
0166             float quality = (v2t_ite->val)[v_ite].second;
0167
0168             if ((vtxref == firstVertexRef) && (quality >= input_MinQuality_)) {
0169               v2t_firstvertextracks->push_back(*AMtrkref);
0170             }
0171           }
0172         }
0173       }
0174     }
0175
0176     iEvent.put(std::move(v2t_firstvertextracks), "V2T");
0177   }
0178 }
0179
0180 bool PF_PU_FirstVertexTracks::trackMatch(const Track& track1, const Track& track2) const {
0181   return ((track1).eta() == (track2).eta() && (track1).phi() == (track2).phi() && (track1).chi2() == (track2).chi2() &&
0182           (track1).ndof() == (track2).ndof() && (track1).p() == (track2).p());
0183 }
0184
0185 // ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
0186 void PF_PU_FirstVertexTracks::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0187   //The following says we do not know what parameters are allowed so do no validation
0188   // Please change this to state exactly what you do use, even if it is no parameters
0189   edm::ParameterSetDescription desc;
0190
0191   using namespace edm;
0192   desc.add<InputTag>("AssociationType");
0193   desc.add<InputTag>("AssociationMap");
0194   desc.add<InputTag>("TrackCollection");
0195   desc.add<InputTag>("VertexCollection");
0196   desc.add<int>("MinQuality");
0197
0198   descriptions.addDefault(desc);
0199 }
0200
0201 //define this as a plug-in
0202 DEFINE_FWK_MODULE(PF_PU_FirstVertexTracks);