FinalTrackSelectors/plugins/DuplicateListMerger.cc

0001 /** \class DuplicateListMerger
0002  *
0003  * merges list of merge duplicate tracks with its parent list
0004  *
0005  * \author Matthew Walker
0006  */
0007
0008 #include "RecoTracker/FinalTrackSelectors/interface/TrackCollectionCloner.h"
0009
0010 #include "FWCore/Framework/interface/global/EDProducer.h"
0011 #include "FWCore/Framework/interface/EventSetup.h"
0012 #include "FWCore/Framework/interface/ESHandle.h"
0013 #include "FWCore/Utilities/interface/InputTag.h"
0014 #include "FWCore/Utilities/interface/ESGetToken.h"
0015
0016 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0017 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0018 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0019
0020 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0021 #include "DataFormats/TrackReco/interface/Track.h"
0022
0023 #include "DataFormats/TrackCandidate/interface/TrackCandidate.h"
0024
0025 #include "RecoTracker/FinalTrackSelectors/interface/TrackAlgoPriorityOrder.h"
0026 #include "RecoTracker/Record/interface/CkfComponentsRecord.h"
0027
0028 #include <vector>
0029 #include <algorithm>
0030 #include <string>
0031 #include <iostream>
0032 #include <memory>
0033
0034 // #include "TMVA/Reader.h"
0035
0036 using namespace reco;
0037
0038 namespace {
0039   class DuplicateListMerger final : public edm::global::EDProducer<> {
0040   public:
0041     /// constructor
0042     explicit DuplicateListMerger(const edm::ParameterSet& iPara);
0043     /// destructor
0044     ~DuplicateListMerger() override;
0045
0046     /// alias for container of candidate and input tracks
0047     using CandidateToDuplicate = std::vector<std::pair<int, int>>;
0048
0049     using RecHitContainer = edm::OwnVector<TrackingRecHit>;
0050
0051     using MVACollection = std::vector<float>;
0052     using QualityMaskCollection = std::vector<unsigned char>;
0053
0054     static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0055
0056   private:
0057     /// produce one event
0058     void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
0059
0060   private:
0061     TrackCollectionCloner collectionCloner;
0062     TrackCollectionCloner::Tokens mergedTrackSource_;
0063     TrackCollectionCloner::Tokens originalTrackSource_;
0064
0065     edm::EDGetTokenT<CandidateToDuplicate> candidateComponents_;
0066     edm::EDGetTokenT<std::vector<TrackCandidate>> candidateSource_;
0067
0068     edm::EDGetTokenT<MVACollection> originalMVAValsToken_;
0069     edm::EDGetTokenT<MVACollection> mergedMVAValsToken_;
0070     edm::ESGetToken<TrackAlgoPriorityOrder, CkfComponentsRecord> priorityOrderToken_;
0071
0072     std::string priorityName_;
0073
0074     int diffHitsCut_;
0075   };
0076 }  // namespace
0077
0078 #include "TrackingTools/PatternTools/interface/Trajectory.h"
0079 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0080 #include "DataFormats/TrajectorySeed/interface/TrajectorySeedCollection.h"
0081 #include "TrackingTools/PatternTools/interface/ClusterRemovalRefSetter.h"
0082
0083 #include "FWCore/Framework/interface/Event.h"
0084 #include "CommonTools/Statistics/interface/ChiSquaredProbability.h"
0085
0086 #include <tuple>
0087 #include <array>
0088 #include "CommonTools/Utils/interface/DynArray.h"
0089
0090 namespace {
0091   void DuplicateListMerger::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0092     edm::ParameterSetDescription desc;
0093     desc.add<edm::InputTag>("mergedSource", edm::InputTag());
0094     desc.add<edm::InputTag>("originalSource", edm::InputTag());
0095     desc.add<edm::InputTag>("mergedMVAVals", edm::InputTag());
0096     desc.add<edm::InputTag>("originalMVAVals", edm::InputTag());
0097     desc.add<edm::InputTag>("candidateSource", edm::InputTag());
0098     desc.add<edm::InputTag>("candidateComponents", edm::InputTag());
0099     desc.add<std::string>("trackAlgoPriorityOrder", "trackAlgoPriorityOrder");
0100     desc.add<int>("diffHitsCut", 5);
0101     TrackCollectionCloner::fill(desc);
0102     descriptions.add("DuplicateListMerger", desc);
0103   }
0104
0105   DuplicateListMerger::DuplicateListMerger(const edm::ParameterSet& iPara)
0106       : collectionCloner(producesCollector(), iPara, true),
0107         mergedTrackSource_(iPara.getParameter<edm::InputTag>("mergedSource"), consumesCollector()),
0108         originalTrackSource_(iPara.getParameter<edm::InputTag>("originalSource"), consumesCollector()),
0109         priorityName_(iPara.getParameter<std::string>("trackAlgoPriorityOrder")) {
0110     diffHitsCut_ = iPara.getParameter<int>("diffHitsCut");
0111     candidateSource_ = consumes<std::vector<TrackCandidate>>(iPara.getParameter<edm::InputTag>("candidateSource"));
0112     candidateComponents_ = consumes<CandidateToDuplicate>(iPara.getParameter<edm::InputTag>("candidateComponents"));
0113
0114     mergedMVAValsToken_ = consumes<MVACollection>(iPara.getParameter<edm::InputTag>("mergedMVAVals"));
0115     originalMVAValsToken_ = consumes<MVACollection>(iPara.getParameter<edm::InputTag>("originalMVAVals"));
0116     priorityOrderToken_ = esConsumes<TrackAlgoPriorityOrder, CkfComponentsRecord>(edm::ESInputTag("", priorityName_));
0117
0118     produces<MVACollection>("MVAValues");
0119     produces<QualityMaskCollection>("QualityMasks");
0120   }
0121
0122   DuplicateListMerger::~DuplicateListMerger() { /* no op */
0123   }
0124
0125   void DuplicateListMerger::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0126     TrackCollectionCloner::Producer producer(iEvent, collectionCloner);
0127
0128     auto const& originals = originalTrackSource_.tracks(iEvent);
0129     auto const& merged = mergedTrackSource_.tracks(iEvent);
0130     auto const& candIndices = mergedTrackSource_.indicesInput(iEvent);
0131
0132     edm::Handle<std::vector<TrackCandidate>> candidateH;
0133     iEvent.getByToken(candidateSource_, candidateH);
0134     auto const& candidates = *candidateH;
0135
0136     edm::Handle<CandidateToDuplicate> candidateComponentsH;
0137     iEvent.getByToken(candidateComponents_, candidateComponentsH);
0138     auto const& candidateComponents = *candidateComponentsH;
0139
0140     edm::Handle<MVACollection> originalMVAStore;
0141     edm::Handle<MVACollection> mergedMVAStore;
0142
0143     iEvent.getByToken(originalMVAValsToken_, originalMVAStore);
0144     iEvent.getByToken(mergedMVAValsToken_, mergedMVAStore);
0145
0146     edm::ESHandle<TrackAlgoPriorityOrder> priorityH = iSetup.getHandle(priorityOrderToken_);
0147     auto const& trackAlgoPriorityOrder = *priorityH;
0148
0149     MVACollection mvaVec;
0150
0151     auto mergedMVA = *mergedMVAStore;
0152
0153     //match new tracks to their candidates
0154     std::vector<std::array<int, 3>> matches;
0155     for (int i = 0; i < (int)merged.size(); ++i) {
0156       auto cInd = candIndices[i];
0157       auto const& cand = candidates[cInd];
0158       const reco::Track& matchedTrack = merged[i];
0159
0160       if (mergedMVA[i] < -0.7f)
0161         continue;  // at least "loose"  ( FIXME: take cut value from CutSelector)
0162
0163       // if( ChiSquaredProbability(matchedTrack.chi2(),matchedTrack.ndof()) < minTrkProbCut_)continue;
0164       int dHits = cand.nRecHits() - matchedTrack.recHitsSize();
0165       if (dHits > diffHitsCut_)
0166         continue;
0167       matches.push_back(std::array<int, 3>{{i, candidateComponents[cInd].first, candidateComponents[cInd].second}});
0168     }
0169
0170     //check for candidates/tracks that share merged tracks, select minimum chi2, remove the rest
0171     if (matches.size() > 1)
0172       for (auto matchIter0 = matches.begin(); matchIter0 != matches.end() - 1; ++matchIter0) {
0173         if ((*matchIter0)[0] < 0)
0174           continue;
0175         auto nchi2 = merged[(*matchIter0)[0]].normalizedChi2();
0176         for (auto matchIter1 = matchIter0 + 1; matchIter1 != matches.end(); ++matchIter1) {
0177           if ((*matchIter1)[0] < 0)
0178             continue;
0179           if ((*matchIter0)[1] == (*matchIter1)[1] || (*matchIter0)[1] == (*matchIter1)[2] ||
0180               (*matchIter0)[2] == (*matchIter1)[1] || (*matchIter0)[2] == (*matchIter1)[2]) {
0181             auto nchi2_1 = merged[(*matchIter1)[0]].normalizedChi2();
0182             if (nchi2_1 < nchi2) {
0183               (*matchIter0)[0] = -1;
0184               break;
0185             } else {
0186               (*matchIter1)[0] = -1;
0187             }
0188           }
0189         }
0190       }
0191
0192     // products
0193     auto pmvas = std::make_unique<MVACollection>();
0194     auto pquals = std::make_unique<QualityMaskCollection>();
0195
0196     //add the good merged tracks to the output list, remove input tracks
0197     std::vector<int> inputTracks;
0198
0199     std::vector<unsigned int> selId;
0200     auto ntotTk = matches.size();
0201     // declareDynArray(reco::TrackBase::TrackAlgorithm, ntotTk, algo);
0202     declareDynArray(reco::TrackBase::TrackAlgorithm, ntotTk, oriAlgo);
0203     declareDynArray(reco::TrackBase::AlgoMask, ntotTk, algoMask);
0204
0205     auto nsel = 0U;
0206     for (auto matchIter0 = matches.begin(); matchIter0 != matches.end(); matchIter0++) {
0207       if ((*matchIter0)[0] < 0)
0208         continue;
0209       selId.push_back((*matchIter0)[0]);
0210
0211       pmvas->push_back(mergedMVA[(*matchIter0)[0]]);
0212
0213       const reco::Track& inTrk1 = originals[(*matchIter0)[1]];
0214       const reco::Track& inTrk2 = originals[(*matchIter0)[2]];
0215       oriAlgo[nsel] = std::min(
0216           inTrk1.algo(), inTrk2.algo(), [&](reco::TrackBase::TrackAlgorithm a, reco::TrackBase::TrackAlgorithm b) {
0217             return trackAlgoPriorityOrder.priority(a) < trackAlgoPriorityOrder.priority(b);
0218           });
0219
0220       algoMask[nsel] = inTrk1.algoMask() | inTrk2.algoMask();
0221
0222       pquals->push_back((inTrk1.qualityMask() | inTrk2.qualityMask()));
0223       pquals->back() |= (1 << reco::TrackBase::confirmed);
0224
0225       inputTracks.push_back((*matchIter0)[1]);
0226       inputTracks.push_back((*matchIter0)[2]);
0227
0228       ++nsel;
0229     }
0230
0231     producer(mergedTrackSource_, selId);
0232     assert(producer.selTracks_->size() == pquals->size());
0233
0234     for (auto isel = 0U; isel < nsel; ++isel) {
0235       algoMask[isel].set(reco::TrackBase::duplicateMerge);
0236       auto& otk = (*producer.selTracks_)[isel];
0237       otk.setQualityMask((*pquals)[isel]);
0238       otk.setAlgorithm(reco::TrackBase::duplicateMerge);
0239       otk.setOriginalAlgorithm(oriAlgo[isel]);
0240       otk.setAlgoMask(algoMask[isel]);
0241     }
0242
0243     selId.clear();
0244     for (int i = 0; i < (int)originals.size(); i++) {
0245       const reco::Track& origTrack = originals[i];
0246       if (std::find(inputTracks.begin(), inputTracks.end(), i) != inputTracks.end())
0247         continue;
0248       selId.push_back(i);
0249       pmvas->push_back((*originalMVAStore)[i]);
0250       pquals->push_back(origTrack.qualityMask());
0251     }
0252
0253     producer(originalTrackSource_, selId);
0254     assert(producer.selTracks_->size() == pquals->size());
0255
0256     iEvent.put(std::move(pmvas), "MVAValues");
0257     iEvent.put(std::move(pquals), "QualityMasks");
0258   }
0259
0260 }  // namespace
0261
0262 #include "FWCore/PluginManager/interface/ModuleDef.h"
0263 #include "FWCore/Framework/interface/MakerMacros.h"
0264
0265 DEFINE_FWK_MODULE(DuplicateListMerger);