RecoUtils/src/PF_PU_AssoMapAlgos.cc

0001
0002 #include "CommonTools/RecoUtils/interface/PF_PU_AssoMapAlgos.h"
0003
0004 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0005
0006 #include "DataFormats/Math/interface/deltaR.h"
0007
0008 #include "DataFormats/GeometryCommonDetAlgo/interface/Measurement1D.h"
0009 #include "RecoVertex/VertexTools/interface/VertexDistance3D.h"
0010 #include "RecoVertex/VertexPrimitives/interface/ConvertToFromReco.h"
0011 #include "RecoVertex/VertexPrimitives/interface/VertexState.h"
0012
0013 #include "DataFormats/RecoCandidate/interface/RecoChargedCandidate.h"
0014
0015 #include "TrackingTools/IPTools/interface/IPTools.h"
0016 #include "CommonTools/Egamma/interface/ConversionTools.h"
0017 #include "Geometry/Records/interface/GlobalTrackingGeometryRecord.h"
0018
0019 using namespace edm;
0020 using namespace std;
0021 using namespace reco;
0022
0023 /*************************************************************************************/
0024 /* dedicated constructor for the algorithms                                          */
0025 /*************************************************************************************/
0026
0027 PF_PU_AssoMapAlgos::PF_PU_AssoMapAlgos(const edm::ParameterSet& iConfig, edm::ConsumesCollector& iC)
0028     : token_bField_(iC.esConsumes()),
0029       token_TrackingGeometry_(iC.esConsumes()),
0030       maxNumWarnings_(3),
0031       numWarnings_(0)
0032
0033 {
0034   input_MaxNumAssociations_ = iConfig.getParameter<int>("MaxNumberOfAssociations");
0035
0036   token_VertexCollection_ = iC.consumes<VertexCollection>(iConfig.getParameter<InputTag>("VertexCollection"));
0037
0038   token_BeamSpot_ = iC.consumes<BeamSpot>(iConfig.getParameter<InputTag>("BeamSpot"));
0039
0040   input_doReassociation_ = iConfig.getParameter<bool>("doReassociation");
0041   cleanedColls_ = iConfig.getParameter<bool>("GetCleanedCollections");
0042
0043   ConversionsCollectionToken_ =
0044       iC.consumes<ConversionCollection>(iConfig.getParameter<InputTag>("ConversionsCollection"));
0045
0046   KshortCollectionToken_ =
0047       iC.consumes<VertexCompositeCandidateCollection>(iConfig.getParameter<InputTag>("V0KshortCollection"));
0048   LambdaCollectionToken_ =
0049       iC.consumes<VertexCompositeCandidateCollection>(iConfig.getParameter<InputTag>("V0LambdaCollection"));
0050
0051   NIVertexCollectionToken_ =
0052       iC.consumes<PFDisplacedVertexCollection>(iConfig.getParameter<InputTag>("NIVertexCollection"));
0053
0054   input_FinalAssociation_ = iConfig.getUntrackedParameter<int>("FinalAssociation", 0);
0055
0056   ignoremissingpfcollection_ = iConfig.getParameter<bool>("ignoreMissingCollection");
0057
0058   input_nTrack_ = iConfig.getParameter<double>("nTrackWeight");
0059 }
0060
0061 /*************************************************************************************/
0062 /* get all needed collections at the beginning                                       */
0063 /*************************************************************************************/
0064
0065 void PF_PU_AssoMapAlgos::GetInputCollections(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0066   //get the offline beam spot
0067   iEvent.getByToken(token_BeamSpot_, beamspotH);
0068
0069   //get the conversion collection for the gamma conversions
0070   iEvent.getByToken(ConversionsCollectionToken_, convCollH);
0071   cleanedConvCollP = PF_PU_AssoMapAlgos::GetCleanedConversions(convCollH, beamspotH, cleanedColls_);
0072
0073   //get the vertex composite candidate collection for the Kshort's
0074   iEvent.getByToken(KshortCollectionToken_, vertCompCandCollKshortH);
0075   cleanedKshortCollP = PF_PU_AssoMapAlgos::GetCleanedKshort(vertCompCandCollKshortH, beamspotH, cleanedColls_);
0076
0077   //get the vertex composite candidate collection for the Lambda's
0078   iEvent.getByToken(LambdaCollectionToken_, vertCompCandCollLambdaH);
0079   cleanedLambdaCollP = PF_PU_AssoMapAlgos::GetCleanedLambda(vertCompCandCollLambdaH, beamspotH, cleanedColls_);
0080
0081   //get the displaced vertex collection for nuclear interactions
0082   //create a new bool, true if no displaced vertex collection is in the event, mostly for AOD
0083   missingColls = false;
0084   if (!iEvent.getByToken(NIVertexCollectionToken_, displVertexCollH)) {
0085     if (ignoremissingpfcollection_) {
0086       missingColls = true;
0087
0088       if (numWarnings_ < maxNumWarnings_) {
0089         LogWarning("PF_PU_AssoMapAlgos::GetInputCollections")
0090             << "No Extra objects available in input file --> skipping reconstruction of displaced vertices !!" << endl;
0091         ++numWarnings_;
0092       }
0093     }
0094   } else {
0095     cleanedNICollP = PF_PU_AssoMapAlgos::GetCleanedNI(displVertexCollH, beamspotH, true);
0096   }
0097
0098   //get the input vertex collection
0099   iEvent.getByToken(token_VertexCollection_, vtxcollH);
0100
0101   bFieldH = iSetup.getHandle(token_bField_);
0102   trackingGeometryH = iSetup.getHandle(token_TrackingGeometry_);
0103 }
0104
0105 /*************************************************************************************/
0106 /* create the track-to-vertex and vertex-to-track maps in one go                     */
0107 /*************************************************************************************/
0108 std::pair<std::unique_ptr<TrackToVertexAssMap>, std::unique_ptr<VertexToTrackAssMap>>
0109 PF_PU_AssoMapAlgos::createMappings(edm::Handle<reco::TrackCollection> trkcollH) {
0110   unique_ptr<TrackToVertexAssMap> track2vertex(new TrackToVertexAssMap(vtxcollH, trkcollH));
0111   unique_ptr<VertexToTrackAssMap> vertex2track(new VertexToTrackAssMap(trkcollH, vtxcollH));
0112
0113   int num_vertices = vtxcollH->size();
0114   if (num_vertices < input_MaxNumAssociations_)
0115     input_MaxNumAssociations_ = num_vertices;
0116   vector<VertexRef> vtxColl_help;
0117   if (input_MaxNumAssociations_ == 1)
0118     vtxColl_help = CreateVertexVector(vtxcollH);
0119
0120   //loop over all tracks of the track collection
0121   for (size_t idxTrack = 0; idxTrack < trkcollH->size(); ++idxTrack) {
0122     TrackRef trackref = TrackRef(trkcollH, idxTrack);
0123
0124     TransientTrack transtrk(trackref, &(*bFieldH));
0125     transtrk.setBeamSpot(*beamspotH);
0126     transtrk.setTrackingGeometry(trackingGeometryH);
0127
0128     if (input_MaxNumAssociations_ > 1)
0129       vtxColl_help = CreateVertexVector(vtxcollH);
0130
0131     for (int assoc_ite = 0; assoc_ite < input_MaxNumAssociations_; ++assoc_ite) {
0132       VertexStepPair assocVtx =
0133           FindAssociation(trackref, vtxColl_help, bFieldH, trackingGeometryH, beamspotH, assoc_ite);
0134       int step = assocVtx.second;
0135       double distance = (IPTools::absoluteImpactParameter3D(transtrk, *(assocVtx.first))).second.value();
0136
0137       int quality = DefineQuality(assoc_ite, step, distance);
0138
0139       //std::cout << "associating track: Pt = " << trackref->pt() << ","
0140       //          << " eta = " << trackref->eta() << ", phi = " << trackref->phi()
0141       //          << " to vertex: z = " << associatedVertex.first->position().z() << " with quality q = " << quality << std::endl;
0142
0143       // Insert the best vertex and the pair of track and the quality of this association in the map
0144       track2vertex->insert(assocVtx.first, make_pair(trackref, quality));
0145       vertex2track->insert(trackref, make_pair(assocVtx.first, quality));
0146
0147       //cleanup only if multiple iterations are made
0148       if (input_MaxNumAssociations_ > 1)
0149         PF_PU_AssoMapAlgos::EraseVertex(vtxColl_help, assocVtx.first);
0150     }
0151   }
0152
0153   return {std::move(track2vertex), std::move(vertex2track)};
0154 }
0155
0156 /*************************************************************************************/
0157 /* create the track to vertex association map                                        */
0158 /*************************************************************************************/
0159
0160 std::unique_ptr<TrackToVertexAssMap> PF_PU_AssoMapAlgos::CreateTrackToVertexMap(
0161     edm::Handle<reco::TrackCollection> trkcollH) {
0162   return createMappings(trkcollH).first;
0163 }
0164
0165 /*************************************************************************************/
0166 /* create the vertex to track association map                                        */
0167 /*************************************************************************************/
0168
0169 std::unique_ptr<VertexToTrackAssMap> PF_PU_AssoMapAlgos::CreateVertexToTrackMap(
0170     edm::Handle<reco::TrackCollection> trkcollH) {
0171   return createMappings(trkcollH).second;
0172 }
0173
0174 /*****************************************************************************************/
0175 /* function to sort the vertices in the AssociationMap by the sum of (pT - pT_Error)**2  */
0176 /*****************************************************************************************/
0177
0178 unique_ptr<TrackToVertexAssMap> PF_PU_AssoMapAlgos::SortAssociationMap(TrackToVertexAssMap* trackvertexassInput,
0179                                                                        edm::Handle<reco::TrackCollection> trkcollH) {
0180   //create a new TrackVertexAssMap for the Output which will be sorted
0181   unique_ptr<TrackToVertexAssMap> trackvertexassOutput(new TrackToVertexAssMap(vtxcollH, trkcollH));
0182
0183   //Create and fill a vector of pairs of vertex and the summed (pT-pT_Error)**2 of the tracks associated to the vertex
0184   VertexPtsumVector vertexptsumvector;
0185
0186   //loop over all vertices in the association map
0187   for (TrackToVertexAssMap::const_iterator assomap_ite = trackvertexassInput->begin();
0188        assomap_ite != trackvertexassInput->end();
0189        assomap_ite++) {
0190     const VertexRef assomap_vertexref = assomap_ite->key;
0191     const TrackQualityPairVector trckcoll = assomap_ite->val;
0192
0193     float ptsum = 0;
0194
0195     TrackRef trackref;
0196
0197     //get the tracks associated to the vertex and calculate the manipulated pT**2
0198     for (unsigned int trckcoll_ite = 0; trckcoll_ite < trckcoll.size(); trckcoll_ite++) {
0199       trackref = trckcoll[trckcoll_ite].first;
0200       int quality = trckcoll[trckcoll_ite].second;
0201
0202       if (quality <= 2)
0203         continue;
0204
0205       double man_pT = trackref->pt() - trackref->ptError();
0206       if (man_pT > 0.)
0207         ptsum += man_pT * man_pT;
0208     }
0209
0210     vertexptsumvector.push_back(make_pair(assomap_vertexref, ptsum));
0211   }
0212
0213   while (!vertexptsumvector.empty()) {
0214     VertexRef vertexref_highestpT;
0215     float highestpT = 0.;
0216     int highestpT_index = 0;
0217
0218     for (unsigned int vtxptsumvec_ite = 0; vtxptsumvec_ite < vertexptsumvector.size(); vtxptsumvec_ite++) {
0219       if (vertexptsumvector[vtxptsumvec_ite].second > highestpT) {
0220         vertexref_highestpT = vertexptsumvector[vtxptsumvec_ite].first;
0221         highestpT = vertexptsumvector[vtxptsumvec_ite].second;
0222         highestpT_index = vtxptsumvec_ite;
0223       }
0224     }
0225
0226     //loop over all vertices in the association map
0227     for (TrackToVertexAssMap::const_iterator assomap_ite = trackvertexassInput->begin();
0228          assomap_ite != trackvertexassInput->end();
0229          assomap_ite++) {
0230       const VertexRef assomap_vertexref = assomap_ite->key;
0231       const TrackQualityPairVector trckcoll = assomap_ite->val;
0232
0233       //if the vertex from the association map the vertex with the highest manipulated pT
0234       //insert all associated tracks in the output Association Map
0235       if (assomap_vertexref == vertexref_highestpT)
0236         for (unsigned int trckcoll_ite = 0; trckcoll_ite < trckcoll.size(); trckcoll_ite++)
0237           trackvertexassOutput->insert(assomap_vertexref, trckcoll[trckcoll_ite]);
0238     }
0239
0240     vertexptsumvector.erase(vertexptsumvector.begin() + highestpT_index);
0241   }
0242
0243   return trackvertexassOutput;
0244 }
0245
0246 /********************/
0247 /*                  */
0248 /* Member Functions */
0249 /*                  */
0250 /********************/
0251
0252 /*************************************************************************************/
0253 /* create helping vertex vector to remove associated vertices                        */
0254 /*************************************************************************************/
0255
0256 std::vector<reco::VertexRef> PF_PU_AssoMapAlgos::CreateVertexVector(edm::Handle<reco::VertexCollection> vtxcollH) {
0257   vector<VertexRef> output;
0258   output.reserve(vtxcollH->size());
0259   auto const nVtx = vtxcollH->size();
0260   for (unsigned int i = 0; i < nVtx; ++i)
0261     output.emplace_back(vtxcollH, i);
0262
0263   return output;
0264 }
0265
0266 /****************************************************************************/
0267 /* erase one vertex from the vertex vector                                  */
0268 /****************************************************************************/
0269
0270 void PF_PU_AssoMapAlgos::EraseVertex(std::vector<reco::VertexRef>& vtxcollV, reco::VertexRef toErase) {
0271   for (unsigned int index_vtx = 0; index_vtx < vtxcollV.size(); ++index_vtx) {
0272     VertexRef vertexref = vtxcollV.at(index_vtx);
0273
0274     if (vertexref == toErase) {
0275       vtxcollV.erase(vtxcollV.begin() + index_vtx);
0276       break;
0277     }
0278   }
0279 }
0280
0281 /*************************************************************************************/
0282 /* function to find the closest vertex in 3D for a certain track                     */
0283 /*************************************************************************************/
0284
0285 VertexRef PF_PU_AssoMapAlgos::FindClosestZ(const reco::TrackRef trkref,
0286                                            const std::vector<reco::VertexRef>& vtxcollV,
0287                                            double tWeight) {
0288   double ztrack = trkref->vertex().z();
0289
0290   VertexRef foundVertexRef = vtxcollV.at(0);
0291
0292   double dzmin = 1e5;
0293
0294   //loop over all vertices with a good quality in the vertex collection
0295   for (unsigned int index_vtx = 0; index_vtx < vtxcollV.size(); ++index_vtx) {
0296     const VertexRef& vertexref = vtxcollV.at(index_vtx);
0297
0298     double nTracks = sqrt(vertexref->tracksSize());
0299
0300     double z_distance = fabs(ztrack - vertexref->z());
0301
0302     double weightedDistance = z_distance - tWeight * nTracks;
0303
0304     if (weightedDistance < dzmin) {
0305       dzmin = weightedDistance;
0306       foundVertexRef = vertexref;
0307     }
0308   }
0309
0310   return foundVertexRef;
0311 }
0312
0313 /*************************************************************************************/
0314 /* function to find the closest vertex in 3D for a certain track                     */
0315 /*************************************************************************************/
0316
0317 VertexRef PF_PU_AssoMapAlgos::FindClosest3D(TransientTrack transtrk,
0318                                             const std::vector<reco::VertexRef>& vtxcollV,
0319                                             double tWeight) {
0320   VertexRef foundVertexRef = vtxcollV.at(0);
0321
0322   double d3min = 1e5;
0323
0324   //loop over all vertices with a good quality in the vertex collection
0325   for (unsigned int index_vtx = 0; index_vtx < vtxcollV.size(); ++index_vtx) {
0326     const VertexRef& vertexref = vtxcollV.at(index_vtx);
0327
0328     double nTracks = sqrt(vertexref->tracksSize());
0329
0330     double distance = 1e5;
0331     pair<bool, Measurement1D> IpPair = IPTools::absoluteImpactParameter3D(transtrk, *vertexref);
0332
0333     if (IpPair.first)
0334       distance = IpPair.second.value();
0335
0336     double weightedDistance = distance - tWeight * nTracks;
0337
0338     if (weightedDistance < d3min) {
0339       d3min = weightedDistance;
0340       foundVertexRef = vertexref;
0341     }
0342   }
0343
0344   return foundVertexRef;
0345 }
0346
0347 /****************************************************************************************/
0348 /* function to calculate the deltaR between a vector and a vector connecting two points */
0349 /****************************************************************************************/
0350
0351 double PF_PU_AssoMapAlgos::dR(const math::XYZPoint& vtx_pos,
0352                               const math::XYZVector& vtx_mom,
0353                               edm::Handle<reco::BeamSpot> bsH) {
0354   double bs_x = bsH->x0();
0355   double bs_y = bsH->y0();
0356   double bs_z = bsH->z0();
0357
0358   double connVec_x = vtx_pos.x() - bs_x;
0359   double connVec_y = vtx_pos.y() - bs_y;
0360   double connVec_z = vtx_pos.z() - bs_z;
0361
0362   double connVec_r = sqrt(connVec_x * connVec_x + connVec_y * connVec_y + connVec_z * connVec_z);
0363   double connVec_theta = acos(connVec_z * 1. / connVec_r);
0364
0365   double connVec_eta = -1. * log(tan(connVec_theta * 1. / 2.));
0366   double connVec_phi = atan2(connVec_y, connVec_x);
0367
0368   return deltaR(vtx_mom.eta(), vtx_mom.phi(), connVec_eta, connVec_phi);
0369 }
0370
0371 /*************************************************************************************/
0372 /* function to filter the conversion collection                                      */
0373 /*************************************************************************************/
0374
0375 unique_ptr<ConversionCollection> PF_PU_AssoMapAlgos::GetCleanedConversions(
0376     edm::Handle<reco::ConversionCollection> convCollH, Handle<BeamSpot> bsH, bool cleanedColl) {
0377   unique_ptr<ConversionCollection> cleanedConvColl(new ConversionCollection());
0378
0379   for (unsigned int convcoll_idx = 0; convcoll_idx < convCollH->size(); convcoll_idx++) {
0380     ConversionRef convref(convCollH, convcoll_idx);
0381
0382     if (!cleanedColl) {
0383       cleanedConvColl->push_back(*convref);
0384       continue;
0385     }
0386
0387     if ((convref->nTracks() == 2) && (fabs(convref->pairInvariantMass()) <= 0.1)) {
0388       cleanedConvColl->push_back(*convref);
0389     }
0390   }
0391
0392   return cleanedConvColl;
0393 }
0394
0395 /*************************************************************************************/
0396 /* function to find out if the track comes from a gamma conversion                   */
0397 /*************************************************************************************/
0398
0399 bool PF_PU_AssoMapAlgos::ComesFromConversion(const TrackRef trackref,
0400                                              const ConversionCollection& cleanedConvColl,
0401                                              Conversion* gamma) {
0402   for (unsigned int convcoll_ite = 0; convcoll_ite < cleanedConvColl.size(); convcoll_ite++) {
0403     if (ConversionTools::matchesConversion(trackref, cleanedConvColl.at(convcoll_ite))) {
0404       *gamma = cleanedConvColl.at(convcoll_ite);
0405       return true;
0406     }
0407   }
0408
0409   return false;
0410 }
0411
0412 /********************************************************************************/
0413 /* function to find the closest vertex for a track from a conversion            */
0414 /********************************************************************************/
0415
0416 VertexRef PF_PU_AssoMapAlgos::FindConversionVertex(const reco::TrackRef trackref,
0417                                                    const reco::Conversion& gamma,
0418                                                    ESHandle<MagneticField> bfH,
0419                                                    edm::ESHandle<GlobalTrackingGeometry> tgH,
0420                                                    edm::Handle<reco::BeamSpot> bsH,
0421                                                    const std::vector<reco::VertexRef>& vtxcollV,
0422                                                    double tWeight) {
0423   math::XYZPoint conv_pos = gamma.conversionVertex().position();
0424
0425   math::XYZVector conv_mom(
0426       gamma.refittedPair4Momentum().x(), gamma.refittedPair4Momentum().y(), gamma.refittedPair4Momentum().z());
0427
0428   Track photon(trackref->chi2(), trackref->ndof(), conv_pos, conv_mom, 0, trackref->covariance());
0429
0430   TransientTrack transpho(photon, &(*bfH));
0431   transpho.setBeamSpot(*bsH);
0432   transpho.setTrackingGeometry(tgH);
0433
0434   return FindClosest3D(transpho, vtxcollV, tWeight);
0435 }
0436
0437 /*************************************************************************************/
0438 /* function to filter the Kshort collection                                          */
0439 /*************************************************************************************/
0440
0441 unique_ptr<VertexCompositeCandidateCollection> PF_PU_AssoMapAlgos::GetCleanedKshort(
0442     Handle<VertexCompositeCandidateCollection> KshortsH, Handle<BeamSpot> bsH, bool cleanedColl) {
0443   unique_ptr<VertexCompositeCandidateCollection> cleanedKaonColl(new VertexCompositeCandidateCollection());
0444
0445   for (unsigned int kscoll_idx = 0; kscoll_idx < KshortsH->size(); kscoll_idx++) {
0446     VertexCompositeCandidateRef ksref(KshortsH, kscoll_idx);
0447
0448     if (!cleanedColl) {
0449       cleanedKaonColl->push_back(*ksref);
0450       continue;
0451     }
0452
0453     VertexDistance3D distanceComputer;
0454
0455     GlobalPoint dec_pos = RecoVertex::convertPos(ksref->vertex());
0456
0457     GlobalError decayVertexError = GlobalError(ksref->vertexCovariance(0, 0),
0458                                                ksref->vertexCovariance(0, 1),
0459                                                ksref->vertexCovariance(1, 1),
0460                                                ksref->vertexCovariance(0, 2),
0461                                                ksref->vertexCovariance(1, 2),
0462                                                ksref->vertexCovariance(2, 2));
0463
0464     math::XYZVector dec_mom(ksref->momentum().x(), ksref->momentum().y(), ksref->momentum().z());
0465
0466     GlobalPoint bsPosition = RecoVertex::convertPos(bsH->position());
0467     GlobalError bsError = RecoVertex::convertError(bsH->covariance3D());
0468
0469     double kaon_significance =
0470         (distanceComputer.distance(VertexState(bsPosition, bsError), VertexState(dec_pos, decayVertexError)))
0471             .significance();
0472
0473     if ((ksref->vertex().rho() >= 3.) && (ksref->vertexNormalizedChi2() <= 3.) &&
0474         (fabs(ksref->mass() - kMass) <= 0.01) && (kaon_significance > 15.) &&
0475         (PF_PU_AssoMapAlgos::dR(ksref->vertex(), dec_mom, bsH) <= 0.3)) {
0476       cleanedKaonColl->push_back(*ksref);
0477     }
0478   }
0479
0480   return cleanedKaonColl;
0481 }
0482
0483 /*************************************************************************************/
0484 /* function to filter the Lambda collection                                          */
0485 /*************************************************************************************/
0486
0487 unique_ptr<VertexCompositeCandidateCollection> PF_PU_AssoMapAlgos::GetCleanedLambda(
0488     Handle<VertexCompositeCandidateCollection> LambdasH, Handle<BeamSpot> bsH, bool cleanedColl) {
0489   unique_ptr<VertexCompositeCandidateCollection> cleanedLambdaColl(new VertexCompositeCandidateCollection());
0490
0491   for (unsigned int lambdacoll_idx = 0; lambdacoll_idx < LambdasH->size(); lambdacoll_idx++) {
0492     VertexCompositeCandidateRef lambdaref(LambdasH, lambdacoll_idx);
0493
0494     if (!cleanedColl) {
0495       cleanedLambdaColl->push_back(*lambdaref);
0496       continue;
0497     }
0498
0499     VertexDistance3D distanceComputer;
0500
0501     GlobalPoint dec_pos = RecoVertex::convertPos(lambdaref->vertex());
0502
0503     GlobalError decayVertexError = GlobalError(lambdaref->vertexCovariance(0, 0),
0504                                                lambdaref->vertexCovariance(0, 1),
0505                                                lambdaref->vertexCovariance(1, 1),
0506                                                lambdaref->vertexCovariance(0, 2),
0507                                                lambdaref->vertexCovariance(1, 2),
0508                                                lambdaref->vertexCovariance(2, 2));
0509
0510     math::XYZVector dec_mom(lambdaref->momentum().x(), lambdaref->momentum().y(), lambdaref->momentum().z());
0511
0512     GlobalPoint bsPosition = RecoVertex::convertPos(bsH->position());
0513     GlobalError bsError = RecoVertex::convertError(bsH->covariance3D());
0514
0515     double lambda_significance =
0516         (distanceComputer.distance(VertexState(bsPosition, bsError), VertexState(dec_pos, decayVertexError)))
0517             .significance();
0518
0519     if ((lambdaref->vertex().rho() >= 3.) && (lambdaref->vertexNormalizedChi2() <= 3.) &&
0520         (fabs(lambdaref->mass() - lamMass) <= 0.005) && (lambda_significance > 15.) &&
0521         (PF_PU_AssoMapAlgos::dR(lambdaref->vertex(), dec_mom, bsH) <= 0.3)) {
0522       cleanedLambdaColl->push_back(*lambdaref);
0523     }
0524   }
0525
0526   return cleanedLambdaColl;
0527 }
0528
0529 /*************************************************************************************/
0530 /* function to find out if the track comes from a V0 decay                           */
0531 /*************************************************************************************/
0532
0533 bool PF_PU_AssoMapAlgos::ComesFromV0Decay(const TrackRef trackref,
0534                                           const VertexCompositeCandidateCollection& cleanedKshort,
0535                                           const VertexCompositeCandidateCollection& cleanedLambda,
0536                                           VertexCompositeCandidate* V0) {
0537   //the part for the reassociation of particles from Kshort decays
0538   for (VertexCompositeCandidateCollection::const_iterator iKS = cleanedKshort.begin(); iKS != cleanedKshort.end();
0539        iKS++) {
0540     const RecoChargedCandidate* dauCand1 = dynamic_cast<const RecoChargedCandidate*>(iKS->daughter(0));
0541     TrackRef dauTk1 = dauCand1->track();
0542     const RecoChargedCandidate* dauCand2 = dynamic_cast<const RecoChargedCandidate*>(iKS->daughter(1));
0543     TrackRef dauTk2 = dauCand2->track();
0544
0545     if ((trackref == dauTk1) || (trackref == dauTk2)) {
0546       *V0 = *iKS;
0547       return true;
0548     }
0549   }
0550
0551   //the part for the reassociation of particles from Lambda decays
0552   for (VertexCompositeCandidateCollection::const_iterator iLambda = cleanedLambda.begin();
0553        iLambda != cleanedLambda.end();
0554        iLambda++) {
0555     const RecoChargedCandidate* dauCand1 = dynamic_cast<const RecoChargedCandidate*>(iLambda->daughter(0));
0556     TrackRef dauTk1 = dauCand1->track();
0557     const RecoChargedCandidate* dauCand2 = dynamic_cast<const RecoChargedCandidate*>(iLambda->daughter(1));
0558     TrackRef dauTk2 = dauCand2->track();
0559
0560     if ((trackref == dauTk1) || (trackref == dauTk2)) {
0561       *V0 = *iLambda;
0562       return true;
0563     }
0564   }
0565
0566   return false;
0567 }
0568
0569 /*************************************************************************************/
0570 /* function to find the closest vertex in z for a track from a V0                    */
0571 /*************************************************************************************/
0572
0573 VertexRef PF_PU_AssoMapAlgos::FindV0Vertex(const TrackRef trackref,
0574                                            const VertexCompositeCandidate& V0_vtx,
0575                                            ESHandle<MagneticField> bFieldH,
0576                                            edm::ESHandle<GlobalTrackingGeometry> trackingGeometryH,
0577                                            Handle<BeamSpot> bsH,
0578                                            const std::vector<reco::VertexRef>& vtxcollV,
0579                                            double tWeight) {
0580   const math::XYZPoint& dec_pos = V0_vtx.vertex();
0581
0582   math::XYZVector dec_mom(V0_vtx.momentum().x(), V0_vtx.momentum().y(), V0_vtx.momentum().z());
0583
0584   Track V0(trackref->chi2(), trackref->ndof(), dec_pos, dec_mom, 0, trackref->covariance());
0585
0586   TransientTrack transV0(V0, &(*bFieldH));
0587   transV0.setBeamSpot(*bsH);
0588   transV0.setTrackingGeometry(trackingGeometryH);
0589
0590   return FindClosest3D(transV0, vtxcollV, tWeight);
0591 }
0592
0593 /*************************************************************************************/
0594 /* function to filter the nuclear interaction collection                             */
0595 /*************************************************************************************/
0596
0597 unique_ptr<PFDisplacedVertexCollection> PF_PU_AssoMapAlgos::GetCleanedNI(Handle<PFDisplacedVertexCollection> NuclIntH,
0598                                                                          Handle<BeamSpot> bsH,
0599                                                                          bool cleanedColl) {
0600   unique_ptr<PFDisplacedVertexCollection> cleanedNIColl(new PFDisplacedVertexCollection());
0601
0602   for (PFDisplacedVertexCollection::const_iterator niref = NuclIntH->begin(); niref != NuclIntH->end(); niref++) {
0603     if ((niref->isFake()) || !(niref->isNucl()))
0604       continue;
0605
0606     if (!cleanedColl) {
0607       cleanedNIColl->push_back(*niref);
0608       continue;
0609     }
0610
0611     VertexDistance3D distanceComputer;
0612
0613     GlobalPoint ni_pos = RecoVertex::convertPos(niref->position());
0614     GlobalError interactionVertexError = RecoVertex::convertError(niref->error());
0615
0616     math::XYZVector ni_mom(niref->primaryMomentum().x(), niref->primaryMomentum().y(), niref->primaryMomentum().z());
0617
0618     GlobalPoint bsPosition = RecoVertex::convertPos(bsH->position());
0619     GlobalError bsError = RecoVertex::convertError(bsH->covariance3D());
0620
0621     double nuclint_significance =
0622         (distanceComputer.distance(VertexState(bsPosition, bsError), VertexState(ni_pos, interactionVertexError)))
0623             .significance();
0624
0625     if ((niref->position().rho() >= 3.) && (nuclint_significance > 15.) &&
0626         (PF_PU_AssoMapAlgos::dR(niref->position(), ni_mom, bsH) <= 0.3)) {
0627       cleanedNIColl->push_back(*niref);
0628     }
0629   }
0630
0631   return cleanedNIColl;
0632 }
0633
0634 /*************************************************************************************/
0635 /* function to find out if the track comes from a nuclear interaction                */
0636 /*************************************************************************************/
0637
0638 bool PF_PU_AssoMapAlgos::ComesFromNI(const TrackRef trackref,
0639                                      const PFDisplacedVertexCollection& cleanedNI,
0640                                      PFDisplacedVertex* displVtx) {
0641   //the part for the reassociation of particles from nuclear interactions
0642   for (PFDisplacedVertexCollection::const_iterator iDisplV = cleanedNI.begin(); iDisplV != cleanedNI.end(); iDisplV++) {
0643     if (iDisplV->trackWeight(trackref) > 1.e-5) {
0644       *displVtx = *iDisplV;
0645       return true;
0646     }
0647   }
0648
0649   return false;
0650 }
0651
0652 /*************************************************************************************/
0653 /* function to find the closest vertex in z for a track from a nuclear interaction   */
0654 /*************************************************************************************/
0655
0656 VertexRef PF_PU_AssoMapAlgos::FindNIVertex(const TrackRef trackref,
0657                                            const PFDisplacedVertex& displVtx,
0658                                            ESHandle<MagneticField> bFieldH,
0659                                            edm::ESHandle<GlobalTrackingGeometry> trackingGeometryH,
0660                                            Handle<BeamSpot> bsH,
0661                                            const std::vector<reco::VertexRef>& vtxcollV,
0662                                            double tWeight) {
0663   TrackCollection refittedTracks = displVtx.refittedTracks();
0664
0665   if ((displVtx.isTherePrimaryTracks()) || (displVtx.isThereMergedTracks())) {
0666     for (TrackCollection::const_iterator trkcoll_ite = refittedTracks.begin(); trkcoll_ite != refittedTracks.end();
0667          trkcoll_ite++) {
0668       const TrackBaseRef retrackbaseref = displVtx.originalTrack(*trkcoll_ite);
0669
0670       if (displVtx.isIncomingTrack(retrackbaseref)) {
0671         VertexRef VOAssociation = TrackWeightAssociation(retrackbaseref, vtxcollV);
0672
0673         if (VOAssociation->trackWeight(retrackbaseref) >= 1.e-5) {
0674           return VOAssociation;
0675         }
0676
0677         TransientTrack transIncom(*retrackbaseref, &(*bFieldH));
0678         transIncom.setBeamSpot(*bsH);
0679         transIncom.setTrackingGeometry(trackingGeometryH);
0680
0681         return FindClosest3D(transIncom, vtxcollV, tWeight);
0682       }
0683     }
0684   }
0685
0686   const math::XYZPoint& ni_pos = displVtx.position();
0687
0688   math::XYZVector ni_mom(
0689       displVtx.primaryMomentum().x(), displVtx.primaryMomentum().y(), displVtx.primaryMomentum().z());
0690
0691   Track incom(trackref->chi2(), trackref->ndof(), ni_pos, ni_mom, 0, trackref->covariance());
0692
0693   TransientTrack transIncom(incom, &(*bFieldH));
0694   transIncom.setBeamSpot(*bsH);
0695   transIncom.setTrackingGeometry(trackingGeometryH);
0696
0697   return FindClosest3D(transIncom, vtxcollV, tWeight);
0698 }
0699
0700 /*************************************************************************************/
0701 /* function to find the vertex with the highest TrackWeight for a certain track      */
0702 /*************************************************************************************/
0703 template <typename TREF>
0704 VertexRef PF_PU_AssoMapAlgos::TrackWeightAssociation(const TREF& trackRef,
0705                                                      const std::vector<reco::VertexRef>& vtxcollV) {
0706   VertexRef bestvertexref = vtxcollV.at(0);
0707   float bestweight = 0.;
0708
0709   //loop over all vertices in the vertex collection
0710   for (auto const& vertexref : vtxcollV) {
0711     //get the most probable vertex for the track
0712     float weight = vertexref->trackWeight(trackRef);
0713     if (weight > bestweight) {
0714       bestweight = weight;
0715       bestvertexref = vertexref;
0716     }
0717   }
0718
0719   return bestvertexref;
0720 }
0721
0722 /*************************************************************************************/
0723 /* find an association for a certain track                                           */
0724 /*************************************************************************************/
0725
0726 VertexStepPair PF_PU_AssoMapAlgos::FindAssociation(const reco::TrackRef& trackref,
0727                                                    const std::vector<reco::VertexRef>& vtxColl,
0728                                                    edm::ESHandle<MagneticField> bfH,
0729                                                    edm::ESHandle<GlobalTrackingGeometry> tgH,
0730                                                    edm::Handle<reco::BeamSpot> bsH,
0731                                                    int assocNum) {
0732   VertexRef foundVertex;
0733
0734   //if it is not the first try of an association jump to the final association
0735   //to avoid multiple (secondary) associations and/or unphysical (primary and secondary) associations
0736   if (assocNum > 0)
0737     goto finalStep;
0738
0739   // Step 1: First round of association:
0740   // Find the vertex with the highest track-to-vertex association weight
0741   foundVertex = TrackWeightAssociation(trackref, vtxColl);
0742
0743   if (foundVertex->trackWeight(trackref) >= 1.e-5) {
0744     return make_pair(foundVertex, 0.);
0745   }
0746
0747   // Step 2: Reassociation
0748   // Second round of association:
0749   // In case no vertex with track-to-vertex association weight > 1.e-5 is found,
0750   // check the track originates from a neutral hadron decay, photon conversion or nuclear interaction
0751
0752   if (input_doReassociation_) {
0753     // Test if the track comes from a photon conversion:
0754     // If so, try to find the vertex of the mother particle
0755     Conversion gamma;
0756     if (ComesFromConversion(trackref, *cleanedConvCollP, &gamma)) {
0757       foundVertex = FindConversionVertex(trackref, gamma, bfH, tgH, bsH, vtxColl, input_nTrack_);
0758       return make_pair(foundVertex, 1.);
0759     }
0760
0761     // Test if the track comes from a Kshort or Lambda decay:
0762     // If so, reassociate the track to the vertex of the V0
0763     VertexCompositeCandidate V0;
0764     if (ComesFromV0Decay(trackref, *cleanedKshortCollP, *cleanedLambdaCollP, &V0)) {
0765       foundVertex = FindV0Vertex(trackref, V0, bfH, tgH, bsH, vtxColl, input_nTrack_);
0766       return make_pair(foundVertex, 1.);
0767     }
0768
0769     if (!missingColls) {
0770       // Test if the track comes from a nuclear interaction:
0771       // If so, reassociate the track to the vertex of the incoming particle
0772       PFDisplacedVertex displVtx;
0773       if (ComesFromNI(trackref, *cleanedNICollP, &displVtx)) {
0774         foundVertex = FindNIVertex(trackref, displVtx, bfH, tgH, bsH, vtxColl, input_nTrack_);
0775         return make_pair(foundVertex, 1.);
0776       }
0777     }
0778   }
0779
0780   // Step 3: Final association
0781   // If no vertex is found with track-to-vertex association weight > 1.e-5
0782   // and no reassociation was done do the final association
0783   // look for the closest vertex in 3D or in z/longitudinal distance
0784   // or associate the track always to the first vertex (default)
0785
0786 finalStep:
0787
0788   switch (input_FinalAssociation_) {
0789     case 1: {
0790       // closest in z
0791       foundVertex = FindClosestZ(trackref, vtxColl, input_nTrack_);
0792       break;
0793     }
0794
0795     case 2: {
0796       // closest in 3D
0797       TransientTrack transtrk(trackref, &(*bfH));
0798       transtrk.setBeamSpot(*bsH);
0799       transtrk.setTrackingGeometry(trackingGeometryH);
0800
0801       foundVertex = FindClosest3D(transtrk, vtxColl, input_nTrack_);
0802       break;
0803     }
0804
0805     default: {
0806       // allways first vertex
0807       foundVertex = vtxColl.at(0);
0808       break;
0809     }
0810   }
0811
0812   return make_pair(foundVertex, 2.);
0813 }
0814
0815 /*************************************************************************************/
0816 /* get the quality for a certain association                                         */
0817 /*************************************************************************************/
0818
0819 int PF_PU_AssoMapAlgos::DefineQuality(int assoc_ite, int step, double distance) {
0820   int quality = 0;
0821
0822   switch (step) {
0823     case 0: {
0824       //TrackWeight association
0825       if (distance <= tw_90) {
0826         quality = 5;
0827       } else {
0828         if (distance <= tw_70) {
0829           quality = 4;
0830         } else {
0831           if (distance <= tw_50) {
0832             quality = 3;
0833           } else {
0834             quality = 2;
0835           }
0836         }
0837       }
0838       break;
0839     }
0840
0841     case 1: {
0842       //Secondary association
0843       if (distance <= sec_70) {
0844         quality = 4;
0845       } else {
0846         if (distance <= sec_50) {
0847           quality = 3;
0848         } else {
0849           quality = 2;
0850         }
0851       }
0852       break;
0853     }
0854
0855     case 2: {
0856       //Final association
0857       if (assoc_ite == 1) {
0858         quality = 1;
0859       } else {
0860         if (assoc_ite >= 2) {
0861           quality = 0;
0862         } else {
0863           if (distance <= fin_70) {
0864             quality = 4;
0865           } else {
0866             if (distance <= fin_50) {
0867               quality = 3;
0868             } else {
0869               quality = 2;
0870             }
0871           }
0872         }
0873       }
0874       break;
0875     }
0876
0877     default: {
0878       quality = -1;
0879       break;
0880     }
0881   }
0882
0883   return quality;
0884 }