CommonAlignmentProducer/src/AlignmentTracksFromVertexSelector.cc

0001 #include "Alignment/CommonAlignmentProducer/interface/AlignmentTracksFromVertexSelector.h"
0002
0003 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0005
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "DataFormats/TrackReco/interface/Track.h"
0008
0009 #include "RecoVertex/VertexPrimitives/interface/TransientVertex.h"
0010 #include "RecoVertex/KalmanVertexFit/interface/KalmanVertexFitter.h"
0011
0012 // vertices
0013 #include "RecoVertex/VertexTools/interface/VertexDistanceXY.h"
0014 #include "RecoVertex/VertexTools/interface/VertexDistance3D.h"
0015
0016 // constructor ----------------------------------------------------------------
0017 AlignmentTrackFromVertexSelector::AlignmentTrackFromVertexSelector(const edm::ParameterSet& cfg,
0018                                                                    edm::ConsumesCollector& iC)
0019     : ttbESToken_(
0020           iC.esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"))),
0021       vertexToken_(iC.consumes<reco::VertexCollection>(cfg.getParameter<edm::InputTag>("vertices"))),
0022       diLeptonToken_(iC.consumes<reco::TrackCollection>(cfg.getParameter<edm::InputTag>("leptonTracks"))),
0023       useClosestVertex_(cfg.getParameter<bool>("useClosestVertexToDilepton")),
0024       vertexIndex_(cfg.getParameter<unsigned int>("vertexIndex")) {}
0025
0026 // destructor -----------------------------------------------------------------
0027 AlignmentTrackFromVertexSelector::~AlignmentTrackFromVertexSelector() {}
0028
0029 // compute the closest vertex to di-lepton ------------------------------------
0030 const reco::Vertex* AlignmentTrackFromVertexSelector::findClosestVertex(const reco::TrackCollection& leptonTracks,
0031                                                                         const reco::VertexCollection* vertices,
0032                                                                         const edm::EventSetup& setup) const {
0033   reco::Vertex* defaultVtx = nullptr;
0034
0035   // fill the transient track collection with the lepton tracks
0036   const TransientTrackBuilder* theB = &setup.getData(ttbESToken_);
0037   std::vector<reco::TransientTrack> tks;
0038   for (const auto& track : leptonTracks) {
0039     reco::TransientTrack trajectory = theB->build(track);
0040     tks.push_back(trajectory);
0041   }
0042
0043   // compute the secondary vertex
0044   TransientVertex aTransVtx;
0045   KalmanVertexFitter kalman(true);
0046   aTransVtx = kalman.vertex(tks);
0047
0048   if (!aTransVtx.isValid())
0049     return defaultVtx;
0050
0051   // find the closest vertex to the secondary vertex in 3D
0052   VertexDistance3D vertTool3D;
0053   float minD = 9999.;
0054   int closestVtxIndex = 0;
0055   int counter = 0;
0056   for (const auto& vtx : *vertices) {
0057     double dist3D = vertTool3D.distance(aTransVtx, vtx).value();
0058     if (dist3D < minD) {
0059       minD = dist3D;
0060       closestVtxIndex = counter;
0061     }
0062     counter++;
0063   }
0064   if ((*vertices).at(closestVtxIndex).isValid()) {
0065     return &(vertices->at(closestVtxIndex));
0066   } else {
0067     return defaultVtx;
0068   }
0069 }
0070
0071 // do selection ---------------------------------------------------------------
0072 AlignmentTrackFromVertexSelector::Tracks AlignmentTrackFromVertexSelector::select(
0073     const edm::Handle<reco::TrackCollection>& tc, const edm::Event& evt, const edm::EventSetup& setup) const {
0074   Tracks result;
0075
0076   std::vector<unsigned int> thePVkeys;
0077
0078   // get collection of reconstructed vertices from event
0079   edm::Handle<reco::VertexCollection> vertexHandle = evt.getHandle(vertexToken_);
0080
0081   // fill the vector of keys
0082   if (vertexHandle.isValid()) {
0083     const reco::VertexCollection* vertices = vertexHandle.product();
0084     const reco::Vertex* theVtx = nullptr;
0085
0086     if (useClosestVertex_) {
0087       // get collection of the di-lepton traxks
0088       const auto& leptonTracks = evt.get(diLeptonToken_);
0089       theVtx = findClosestVertex(leptonTracks, vertices, setup);
0090     } else {
0091       if ((*vertices).at(vertexIndex_).isValid()) {
0092         theVtx = &(vertices->at(vertexIndex_));
0093       }
0094     }
0095
0096     if (theVtx) {
0097       for (auto tv = theVtx->tracks_begin(); tv != theVtx->tracks_end(); tv++) {
0098         if (tv->isNonnull()) {
0099           const reco::TrackRef trackRef = tv->castTo<reco::TrackRef>();
0100           thePVkeys.push_back(trackRef.key());
0101         }
0102       }
0103     }
0104   }
0105
0106   std::sort(thePVkeys.begin(), thePVkeys.end());
0107
0108   LogDebug("AlignmentTrackFromVertexSelector")
0109       << "after collecting the PV keys: thePVkeys.size()" << thePVkeys.size() << std::endl;
0110   for (const auto& key : thePVkeys) {
0111     LogDebug("AlignmentTrackFromVertexSelector") << key << ", ";
0112   }
0113   LogDebug("AlignmentTrackFromVertexSelector") << std::endl;
0114
0115   if (tc.isValid()) {
0116     int indx(0);
0117     // put the track in the collection is it was used for the vertex
0118     for (reco::TrackCollection::const_iterator tk = tc->begin(); tk != tc->end(); ++tk, ++indx) {
0119       reco::TrackRef trackRef = reco::TrackRef(tc, indx);
0120       if (std::find(thePVkeys.begin(), thePVkeys.end(), trackRef.key()) != thePVkeys.end()) {
0121         LogDebug("AlignmentTrackFromVertexSelector") << "track index: " << indx << "filling result vector" << std::endl;
0122         result.push_back(&(*tk));
0123       }  // if a valid key is found
0124     }  // end loop over tracks
0125   }  // if the handle is valid
0126   return result;
0127 }