FinalTrackSelectors/plugins/TrackCollectionMerger.cc

0001 #include "RecoTracker/FinalTrackSelectors/interface/TrackCollectionCloner.h"
0002
0003 #include "FWCore/Framework/interface/Frameworkfwd.h"
0004 #include "FWCore/Framework/interface/global/EDProducer.h"
0005
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "FWCore/Framework/interface/EventSetup.h"
0008 #include "FWCore/Framework/interface/ESHandle.h"
0009 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0010 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0011 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0012
0013 #include "FWCore/Utilities/interface/ESGetToken.h"
0014 #include "FWCore/Utilities/interface/InputTag.h"
0015 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0016 #include "DataFormats/TrackReco/interface/Track.h"
0017 #include "DataFormats/TrackingRecHit/interface/TrackingRecHit.h"
0018 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0019 #include "TrackingTools/PatternTools/interface/Trajectory.h"
0020 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0021
0022 #include "RecoTracker/FinalTrackSelectors/interface/TrackAlgoPriorityOrder.h"
0023 #include "RecoTracker/Record/interface/CkfComponentsRecord.h"
0024
0025 #include <vector>
0026 #include <memory>
0027 #include <cassert>
0028 namespace {
0029   class TrackCollectionMerger final : public edm::global::EDProducer<> {
0030   public:
0031     explicit TrackCollectionMerger(const edm::ParameterSet& conf)
0032         : collectionCloner(producesCollector(), conf, true),
0033           priorityName_(conf.getParameter<std::string>("trackAlgoPriorityOrder")),
0034           m_priorityToken(esConsumes<TrackAlgoPriorityOrder, CkfComponentsRecord>(edm::ESInputTag("", priorityName_))),
0035           m_foundHitBonus(conf.getParameter<double>("foundHitBonus")),
0036           m_lostHitPenalty(conf.getParameter<double>("lostHitPenalty")),
0037           m_shareFrac(conf.getParameter<double>("shareFrac")),
0038           m_minShareHits(conf.getParameter<unsigned int>("minShareHits")),
0039           m_minQuality(reco::TrackBase::qualityByName(conf.getParameter<std::string>("minQuality"))),
0040           m_allowFirstHitShare(conf.getParameter<bool>("allowFirstHitShare")),
0041           m_enableMerging(conf.getParameter<bool>("enableMerging")) {
0042       for (auto const& it : conf.getParameter<std::vector<edm::InputTag>>("trackProducers"))
0043         srcColls.emplace_back(it, consumesCollector());
0044       for (auto const& it : conf.getParameter<std::vector<std::string>>("inputClassifiers")) {
0045         srcMVAs.push_back(consumes<MVACollection>(edm::InputTag(it, "MVAValues")));
0046         srcQuals.push_back(consumes<QualityMaskCollection>(edm::InputTag(it, "QualityMasks")));
0047       }
0048
0049       assert(srcColls.size() == srcQuals.size());
0050
0051       produces<MVACollection>("MVAValues");
0052       produces<QualityMaskCollection>("QualityMasks");
0053     }
0054
0055     static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0056       edm::ParameterSetDescription desc;
0057       desc.add<std::vector<edm::InputTag>>("trackProducers", std::vector<edm::InputTag>());
0058       desc.add<std::vector<std::string>>("inputClassifiers", std::vector<std::string>());
0059       desc.add<std::string>("trackAlgoPriorityOrder", "trackAlgoPriorityOrder");
0060       desc.add<double>("shareFrac", .19);
0061       desc.add<double>("foundHitBonus", 10.);
0062       desc.add<double>("lostHitPenalty", 5.);
0063       desc.add<unsigned int>("minShareHits", 2);
0064       desc.add<bool>("allowFirstHitShare", true);
0065       desc.add<bool>("enableMerging", true);
0066       desc.add<std::string>("minQuality", "loose");
0067       TrackCollectionCloner::fill(desc);
0068       descriptions.add("TrackCollectionMerger", desc);
0069     }
0070
0071   private:
0072     TrackCollectionCloner collectionCloner;
0073     std::vector<TrackCollectionCloner::Tokens> srcColls;
0074
0075     using MVACollection = std::vector<float>;
0076     using QualityMaskCollection = std::vector<unsigned char>;
0077
0078     using IHit = std::pair<unsigned int, TrackingRecHit const*>;
0079     using IHitV = std::vector<IHit>;
0080
0081     std::vector<edm::EDGetTokenT<MVACollection>> srcMVAs;
0082     std::vector<edm::EDGetTokenT<QualityMaskCollection>> srcQuals;
0083
0084     std::string priorityName_;
0085     edm::ESGetToken<TrackAlgoPriorityOrder, CkfComponentsRecord> m_priorityToken;
0086
0087     float m_foundHitBonus;
0088     float m_lostHitPenalty;
0089     float m_shareFrac;
0090     unsigned int m_minShareHits;
0091     reco::TrackBase::TrackQuality m_minQuality;
0092     bool m_allowFirstHitShare;
0093     bool m_enableMerging;
0094
0095     void produce(edm::StreamID, edm::Event& evt, const edm::EventSetup&) const override;
0096
0097     bool areDuplicate(IHitV const& rh1, IHitV const& rh2) const;
0098   };
0099
0100 }  // namespace
0101
0102 #include "CommonTools/Utils/interface/DynArray.h"
0103
0104 namespace {
0105
0106   void TrackCollectionMerger::produce(edm::StreamID, edm::Event& evt, const edm::EventSetup& es) const {
0107     TrackCollectionCloner::Producer producer(evt, collectionCloner);
0108
0109     edm::ESHandle<TrackAlgoPriorityOrder> priorityH = es.getHandle(m_priorityToken);
0110     auto const& trackAlgoPriorityOrder = *priorityH;
0111
0112     // load collections
0113     auto collsSize = srcColls.size();
0114     auto rSize = 0U;
0115     declareDynArray(reco::TrackCollection const*, collsSize, trackColls);
0116     for (auto i = 0U; i < collsSize; ++i) {
0117       trackColls[i] = &srcColls[i].tracks(evt);
0118       rSize += (*trackColls[i]).size();
0119     }
0120
0121     unsigned char qualMask = ~0;
0122     const bool acceptAll = m_minQuality == reco::TrackBase::undefQuality;
0123     if (!acceptAll)
0124       qualMask = 1 << m_minQuality;
0125
0126     // load tracks
0127     initDynArray(unsigned int, collsSize, nGoods, 0);
0128     declareDynArray(float, rSize, mvas);
0129     declareDynArray(unsigned char, rSize, quals);
0130     declareDynArray(unsigned int, rSize, tkInds);
0131     auto k = 0U;
0132     for (auto i = 0U; i < collsSize; ++i) {
0133       auto const& tkColl = *trackColls[i];
0134       auto size = tkColl.size();
0135       edm::Handle<MVACollection> hmva;
0136       evt.getByToken(srcMVAs[i], hmva);
0137       assert((*hmva).size() == size);
0138       edm::Handle<QualityMaskCollection> hqual;
0139       evt.getByToken(srcQuals[i], hqual);
0140       for (auto j = 0U; j < size; ++j) {
0141         if (!(acceptAll || (qualMask & (*hqual)[j])))
0142           continue;
0143         mvas[k] = (*hmva)[j];
0144         quals[k] = (*hqual)[j];
0145         tkInds[k] = j;
0146         ++k;
0147         ++nGoods[i];
0148       }
0149     }
0150
0151     auto ntotTk = k;
0152     std::vector<bool> selected(ntotTk, true);
0153
0154     declareDynArray(reco::TrackBase::TrackAlgorithm, ntotTk, algo);
0155     declareDynArray(reco::TrackBase::TrackAlgorithm, ntotTk, oriAlgo);
0156     declareDynArray(reco::TrackBase::AlgoMask, ntotTk, algoMask);
0157
0158     // merge (if more than one collection...)
0159
0160     auto merger = [&]() -> void {
0161       // load momentum, hits and score
0162       declareDynArray(reco::TrackBase::Vector, ntotTk, mom);
0163       declareDynArray(float, ntotTk, score);
0164       declareDynArray(IHitV, ntotTk, rh1);
0165
0166       k = 0U;
0167       for (auto i = 0U; i < collsSize; ++i) {
0168         auto const& tkColl = *trackColls[i];
0169         for (auto j = 0U; j < nGoods[i]; ++j) {
0170           auto const& track = tkColl[tkInds[k]];
0171           algo[k] = track.algo();
0172           oriAlgo[k] = track.originalAlgo();
0173           algoMask[k] = track.algoMask();
0174           mom[k] = track.isLooper() ? reco::TrackBase::Vector() : track.momentum();
0175           auto validHits = track.numberOfValidHits();
0176           auto validPixelHits = track.hitPattern().numberOfValidPixelHits();
0177           auto lostHits = track.numberOfLostHits();
0178           score[k] = m_foundHitBonus * validPixelHits + m_foundHitBonus * validHits - m_lostHitPenalty * lostHits -
0179                      track.chi2();
0180
0181           auto& rhv = rh1[k];
0182           rhv.reserve(validHits);
0183           auto compById = [](IHit const& h1, IHit const& h2) { return h1.first < h2.first; };
0184           for (auto it = track.recHitsBegin(); it != track.recHitsEnd(); ++it) {
0185             auto const& hit = *(*it);
0186             auto id = hit.rawId();
0187             if LIKELY (hit.isValid()) {
0188               rhv.emplace_back(id, &hit);
0189               std::push_heap(rhv.begin(), rhv.end(), compById);
0190             }
0191           }
0192           std::sort_heap(rhv.begin(), rhv.end(), compById);
0193
0194           ++k;
0195         }
0196       }
0197       assert(ntotTk == k);
0198
0199       auto seti = [&](unsigned int ii, unsigned int jj) {
0200         selected[jj] = false;
0201         selected[ii] = true;
0202         if (trackAlgoPriorityOrder.priority(oriAlgo[jj]) < trackAlgoPriorityOrder.priority(oriAlgo[ii]))
0203           oriAlgo[ii] = oriAlgo[jj];
0204         algoMask[ii] |= algoMask[jj];
0205         quals[ii] |= (1 << reco::TrackBase::confirmed);
0206         algoMask[jj] = algoMask[ii];  // in case we keep discarded
0207         quals[jj] |= (1 << reco::TrackBase::discarded);
0208       };
0209
0210       auto iStart2 = 0U;
0211       for (auto i = 0U; i < collsSize - 1; ++i) {
0212         auto iStart1 = iStart2;
0213         iStart2 = iStart1 + nGoods[i];
0214         for (auto t1 = iStart1; t1 < iStart2; ++t1) {
0215           if (!selected[t1])
0216             continue;
0217           auto score1 = score[t1];
0218           for (auto t2 = iStart2; t2 < ntotTk; ++t2) {
0219             if (!selected[t1])
0220               break;
0221             if (!selected[t2])
0222               continue;
0223             if (mom[t1].Dot(mom[t2]) < 0)
0224               continue;  // do not bother if in opposite hemespheres...
0225             if (!areDuplicate(rh1[t1], rh1[t2]))
0226               continue;
0227             auto score2 = score[t2];
0228
0229             // difference rather than ratio due to possible negative values for score
0230             constexpr float almostSame = 0.01f;
0231             if (score1 - score2 > almostSame) {
0232               seti(t1, t2);
0233             } else if (score2 - score1 > almostSame) {
0234               seti(t2, t1);
0235             } else {  // take best
0236               constexpr unsigned int qmask =
0237                   (1 << reco::TrackBase::loose | 1 << reco::TrackBase::tight | 1 << reco::TrackBase::highPurity);
0238               if ((quals[t1] & qmask) == (quals[t2] & qmask)) {
0239                 // take first
0240                 if (trackAlgoPriorityOrder.priority(algo[t1]) <= trackAlgoPriorityOrder.priority(algo[t2])) {
0241                   seti(t1, t2);
0242                 } else {
0243                   seti(t2, t1);
0244                 }
0245               } else if ((quals[t1] & qmask) > (quals[t2] & qmask))
0246                 seti(t1, t2);
0247               else
0248                 seti(t2, t1);
0249             }  // end ifs...
0250
0251           }  // end t2
0252         }    // end t1
0253       }      // end colls
0254     };       // end merger;
0255
0256     const bool doMerging = m_enableMerging && collsSize > 1;
0257     if (doMerging)
0258       merger();
0259
0260     // products
0261     auto pmvas = std::make_unique<MVACollection>();
0262     auto pquals = std::make_unique<QualityMaskCollection>();
0263
0264     // clone selected tracks...
0265     auto nsel = 0U;
0266     auto iStart2 = 0U;
0267     auto isel = 0U;
0268     for (auto i = 0U; i < collsSize; ++i) {
0269       std::vector<unsigned int> selId;
0270       std::vector<unsigned int> tid;
0271       auto iStart1 = iStart2;
0272       iStart2 = iStart1 + nGoods[i];
0273       assert(producer.selTracks_->size() == isel);
0274       for (auto t1 = iStart1; t1 < iStart2; ++t1) {
0275         if (!selected[t1])
0276           continue;
0277         ++nsel;
0278         tid.push_back(t1);
0279         selId.push_back(tkInds[t1]);
0280         pmvas->push_back(mvas[t1]);
0281         pquals->push_back(quals[t1]);
0282       }
0283       producer(srcColls[i], selId);
0284       assert(producer.selTracks_->size() == nsel);
0285       assert(tid.size() == nsel - isel);
0286       auto k = 0U;
0287       for (; isel < nsel; ++isel) {
0288         auto& otk = (*producer.selTracks_)[isel];
0289         otk.setQualityMask((*pquals)[isel]);  // needed also without merging
0290         if (doMerging) {
0291           otk.setOriginalAlgorithm(oriAlgo[tid[k]]);
0292           otk.setAlgoMask(algoMask[tid[k++]]);
0293         }
0294       }
0295       if (doMerging)
0296         assert(tid.size() == k);
0297     }
0298
0299     assert(producer.selTracks_->size() == pmvas->size());
0300
0301     // std::cout << "TrackCollectionMerger: sel tracks " << producer.selTracks_->size() << std::endl;
0302
0303     evt.put(std::move(pmvas), "MVAValues");
0304     evt.put(std::move(pquals), "QualityMasks");
0305   }
0306
0307   bool TrackCollectionMerger::areDuplicate(IHitV const& rh1, IHitV const& rh2) const {
0308     auto nh1 = rh1.size();
0309     auto nh2 = rh2.size();
0310
0311     auto share = [](const TrackingRecHit* it, const TrackingRecHit* jt) -> bool {
0312       return it->sharesInput(jt, TrackingRecHit::some);
0313     };
0314
0315     //loop over rechits
0316     int noverlap = 0;
0317     int firstoverlap = 0;
0318     // check first hit  (should use REAL first hit?)
0319     if (m_allowFirstHitShare && rh1[0].first == rh2[0].first) {
0320       if (share(rh1[0].second, rh2[0].second))
0321         firstoverlap = 1;
0322     }
0323
0324     // exploit sorting
0325     unsigned int jh = 0;
0326     unsigned int ih = 0;
0327     while (ih != nh1 && jh != nh2) {
0328       auto const id1 = rh1[ih].first;
0329       auto const id2 = rh2[jh].first;
0330       if (id1 < id2)
0331         ++ih;
0332       else if (id2 < id1)
0333         ++jh;
0334       else {
0335         if (share(rh1[ih].second, rh2[jh].second))
0336           noverlap++;
0337         ++jh;
0338         ++ih;
0339       }  // equal ids
0340     }    //loop over ih & jh
0341
0342     return noverlap >= int(m_minShareHits) &&
0343            (noverlap - firstoverlap) > (std::min(nh1, nh2) - firstoverlap) * m_shareFrac;
0344   }
0345
0346 }  // namespace
0347
0348 #include "FWCore/PluginManager/interface/ModuleDef.h"
0349 #include "FWCore/Framework/interface/MakerMacros.h"
0350
0351 DEFINE_FWK_MODULE(TrackCollectionMerger);