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   void DuplicateListMerger::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0125     TrackCollectionCloner::Producer producer(iEvent, collectionCloner);
0126
0127     auto const& originals = originalTrackSource_.tracks(iEvent);
0128     auto const& merged = mergedTrackSource_.tracks(iEvent);
0129     auto const& candIndices = mergedTrackSource_.indicesInput(iEvent);
0130
0131     edm::Handle<std::vector<TrackCandidate>> candidateH;
0132     iEvent.getByToken(candidateSource_, candidateH);
0133     auto const& candidates = *candidateH;
0134
0135     edm::Handle<CandidateToDuplicate> candidateComponentsH;
0136     iEvent.getByToken(candidateComponents_, candidateComponentsH);
0137     auto const& candidateComponents = *candidateComponentsH;
0138
0139     edm::Handle<MVACollection> originalMVAStore;
0140     edm::Handle<MVACollection> mergedMVAStore;
0141
0142     iEvent.getByToken(originalMVAValsToken_, originalMVAStore);
0143     iEvent.getByToken(mergedMVAValsToken_, mergedMVAStore);
0144
0145     edm::ESHandle<TrackAlgoPriorityOrder> priorityH = iSetup.getHandle(priorityOrderToken_);
0146     auto const& trackAlgoPriorityOrder = *priorityH;
0147
0148     MVACollection mvaVec;
0149
0150     auto mergedMVA = *mergedMVAStore;
0151
0152     //match new tracks to their candidates
0153     std::vector<std::array<int, 3>> matches;
0154     for (int i = 0; i < (int)merged.size(); ++i) {
0155       auto cInd = candIndices[i];
0156       auto const& cand = candidates[cInd];
0157       const reco::Track& matchedTrack = merged[i];
0158
0159       if (mergedMVA[i] < -0.7f)
0160         continue;  // at least "loose"  ( FIXME: take cut value from CutSelector)
0161
0162       // if( ChiSquaredProbability(matchedTrack.chi2(),matchedTrack.ndof()) < minTrkProbCut_)continue;
0163       int dHits = cand.nRecHits() - matchedTrack.recHitsSize();
0164       if (dHits > diffHitsCut_)
0165         continue;
0166       matches.push_back(std::array<int, 3>{{i, candidateComponents[cInd].first, candidateComponents[cInd].second}});
0167     }
0168
0169     //check for candidates/tracks that share merged tracks, select minimum chi2, remove the rest
0170     if (matches.size() > 1)
0171       for (auto matchIter0 = matches.begin(); matchIter0 != matches.end() - 1; ++matchIter0) {
0172         if ((*matchIter0)[0] < 0)
0173           continue;
0174         auto nchi2 = merged[(*matchIter0)[0]].normalizedChi2();
0175         for (auto matchIter1 = matchIter0 + 1; matchIter1 != matches.end(); ++matchIter1) {
0176           if ((*matchIter1)[0] < 0)
0177             continue;
0178           if ((*matchIter0)[1] == (*matchIter1)[1] || (*matchIter0)[1] == (*matchIter1)[2] ||
0179               (*matchIter0)[2] == (*matchIter1)[1] || (*matchIter0)[2] == (*matchIter1)[2]) {
0180             auto nchi2_1 = merged[(*matchIter1)[0]].normalizedChi2();
0181             if (nchi2_1 < nchi2) {
0182               (*matchIter0)[0] = -1;
0183               break;
0184             } else {
0185               (*matchIter1)[0] = -1;
0186             }
0187           }
0188         }
0189       }
0190
0191     // products
0192     auto pmvas = std::make_unique<MVACollection>();
0193     auto pquals = std::make_unique<QualityMaskCollection>();
0194
0195     //add the good merged tracks to the output list, remove input tracks
0196     std::vector<int> inputTracks;
0197
0198     std::vector<unsigned int> selId;
0199     auto ntotTk = matches.size();
0200     // declareDynArray(reco::TrackBase::TrackAlgorithm, ntotTk, algo);
0201     declareDynArray(reco::TrackBase::TrackAlgorithm, ntotTk, oriAlgo);
0202     declareDynArray(reco::TrackBase::AlgoMask, ntotTk, algoMask);
0203
0204     auto nsel = 0U;
0205     for (auto matchIter0 = matches.begin(); matchIter0 != matches.end(); matchIter0++) {
0206       if ((*matchIter0)[0] < 0)
0207         continue;
0208       selId.push_back((*matchIter0)[0]);
0209
0210       pmvas->push_back(mergedMVA[(*matchIter0)[0]]);
0211
0212       const reco::Track& inTrk1 = originals[(*matchIter0)[1]];
0213       const reco::Track& inTrk2 = originals[(*matchIter0)[2]];
0214       oriAlgo[nsel] = std::min(
0215           inTrk1.algo(), inTrk2.algo(), [&](reco::TrackBase::TrackAlgorithm a, reco::TrackBase::TrackAlgorithm b) {
0216             return trackAlgoPriorityOrder.priority(a) < trackAlgoPriorityOrder.priority(b);
0217           });
0218
0219       algoMask[nsel] = inTrk1.algoMask() | inTrk2.algoMask();
0220
0221       pquals->push_back((inTrk1.qualityMask() | inTrk2.qualityMask()));
0222       pquals->back() |= (1 << reco::TrackBase::confirmed);
0223
0224       inputTracks.push_back((*matchIter0)[1]);
0225       inputTracks.push_back((*matchIter0)[2]);
0226
0227       ++nsel;
0228     }
0229
0230     producer(mergedTrackSource_, selId);
0231     assert(producer.selTracks_->size() == pquals->size());
0232
0233     for (auto isel = 0U; isel < nsel; ++isel) {
0234       algoMask[isel].set(reco::TrackBase::duplicateMerge);
0235       auto& otk = (*producer.selTracks_)[isel];
0236       otk.setQualityMask((*pquals)[isel]);
0237       otk.setAlgorithm(reco::TrackBase::duplicateMerge);
0238       otk.setOriginalAlgorithm(oriAlgo[isel]);
0239       otk.setAlgoMask(algoMask[isel]);
0240     }
0241
0242     selId.clear();
0243     for (int i = 0; i < (int)originals.size(); i++) {
0244       const reco::Track& origTrack = originals[i];
0245       if (std::find(inputTracks.begin(), inputTracks.end(), i) != inputTracks.end())
0246         continue;
0247       selId.push_back(i);
0248       pmvas->push_back((*originalMVAStore)[i]);
0249       pquals->push_back(origTrack.qualityMask());
0250     }
0251
0252     producer(originalTrackSource_, selId);
0253     assert(producer.selTracks_->size() == pquals->size());
0254
0255     iEvent.put(std::move(pmvas), "MVAValues");
0256     iEvent.put(std::move(pquals), "QualityMasks");
0257   }
0258
0259 }  // namespace
0260
0261 #include "FWCore/PluginManager/interface/ModuleDef.h"
0262 #include "FWCore/Framework/interface/MakerMacros.h"
0263
0264 DEFINE_FWK_MODULE(DuplicateListMerger);