plugins/importers/GeneralTracksImporter.cc

0001 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0002 #include "RecoParticleFlow/PFProducer/interface/BlockElementImporterBase.h"
0003 #include "DataFormats/ParticleFlowReco/interface/PFBlockElementTrack.h"
0004 #include "DataFormats/ParticleFlowReco/interface/PFRecTrack.h"
0005 #include "DataFormats/TrackReco/interface/Track.h"
0006 #include "DataFormats/MuonReco/interface/Muon.h"
0007 #include "DataFormats/Common/interface/ValueMap.h"
0008 #include "DataFormats/HGCalReco/interface/TICLSeedingRegion.h"
0009 #include "RecoParticleFlow/PFProducer/interface/PFMuonAlgo.h"
0010 #include "RecoParticleFlow/PFTracking/interface/PFTrackAlgoTools.h"
0011
0012 class GeneralTracksImporter : public BlockElementImporterBase {
0013 public:
0014   GeneralTracksImporter(const edm::ParameterSet& conf, edm::ConsumesCollector& cc)
0015       : BlockElementImporterBase(conf, cc),
0016         src_(cc.consumes<reco::PFRecTrackCollection>(conf.getParameter<edm::InputTag>("source"))),
0017         vetoEndcap_(conf.getParameter<bool>("vetoEndcap")),
0018         muons_(cc.consumes<reco::MuonCollection>(conf.getParameter<edm::InputTag>("muonSrc"))),
0019         trackQuality_(reco::TrackBase::qualityByName(conf.getParameter<std::string>("trackQuality"))),
0020         DPtovPtCut_(conf.getParameter<std::vector<double>>("DPtOverPtCuts_byTrackAlgo")),
0021         NHitCut_(conf.getParameter<std::vector<unsigned>>("NHitCuts_byTrackAlgo")),
0022         useIterTracking_(conf.getParameter<bool>("useIterativeTracking")),
0023         cleanBadConvBrems_(conf.getParameter<bool>("cleanBadConvertedBrems")),
0024         muonMaxDPtOPt_(conf.getParameter<double>("muonMaxDPtOPt")) {
0025     if (vetoEndcap_) {
0026       vetoMode_ = conf.getParameter<unsigned>("vetoMode");
0027       switch (vetoMode_) {
0028         case pfRecTrackCollection:
0029           vetoPFTracksSrc_ = cc.consumes<reco::PFRecTrackCollection>(conf.getParameter<edm::InputTag>("vetoSrc"));
0030           break;
0031         case ticlSeedingRegion:
0032           vetoTICLSeedingSrc_ =
0033               cc.consumes<std::vector<TICLSeedingRegion>>(conf.getParameter<edm::InputTag>("vetoSrc"));
0034           tracksSrc_ = cc.consumes<reco::TrackCollection>(conf.getParameter<edm::InputTag>("tracksSrc"));
0035           break;
0036         case pfCandidateCollection:
0037           vetoPFCandidatesSrc_ = cc.consumes<reco::PFCandidateCollection>(conf.getParameter<edm::InputTag>("vetoSrc"));
0038           break;
0039       }  // switch
0040     }
0041   }
0042
0043   void importToBlock(const edm::Event&, ElementList&) const override;
0044
0045 private:
0046   const edm::EDGetTokenT<reco::PFRecTrackCollection> src_;
0047   const bool vetoEndcap_;
0048   const edm::EDGetTokenT<reco::MuonCollection> muons_;
0049   const reco::TrackBase::TrackQuality trackQuality_;
0050   const std::vector<double> DPtovPtCut_;
0051   const std::vector<unsigned> NHitCut_;
0052   const bool useIterTracking_, cleanBadConvBrems_;
0053   const double muonMaxDPtOPt_;
0054   unsigned int vetoMode_;
0055   edm::EDGetTokenT<reco::PFRecTrackCollection> vetoPFTracksSrc_;
0056   edm::EDGetTokenT<std::vector<TICLSeedingRegion>> vetoTICLSeedingSrc_;
0057   edm::EDGetTokenT<reco::TrackCollection> tracksSrc_;
0058   edm::EDGetTokenT<reco::PFCandidateCollection> vetoPFCandidatesSrc_;
0059 };
0060
0061 DEFINE_EDM_PLUGIN(BlockElementImporterFactory, GeneralTracksImporter, "GeneralTracksImporter");
0062
0063 void GeneralTracksImporter::importToBlock(const edm::Event& e, BlockElementImporterBase::ElementList& elems) const {
0064   typedef BlockElementImporterBase::ElementList::value_type ElementType;
0065   auto tracks = e.getHandle(src_);
0066
0067   typedef std::pair<edm::ProductID, unsigned> TrackProdIDKey;
0068   std::vector<TrackProdIDKey> vetoed;
0069   if (vetoEndcap_) {
0070     switch (vetoMode_) {
0071       case pfRecTrackCollection: {
0072         const auto& vetoes = e.get(vetoPFTracksSrc_);
0073         for (const auto& veto : vetoes)
0074           vetoed.emplace_back(veto.trackRef().id(), veto.trackRef().key());
0075         break;
0076       }
0077       case ticlSeedingRegion: {
0078         const auto& vetoes = e.get(vetoTICLSeedingSrc_);
0079         auto tracksH = e.getHandle(tracksSrc_);
0080         for (const auto& veto : vetoes) {
0081           assert(veto.collectionID == tracksH.id());
0082           reco::TrackRef trkref = reco::TrackRef(tracksH, veto.index);
0083           vetoed.emplace_back(tracksH.id(), veto.index);  // track prod id and key
0084         }
0085         break;
0086       }
0087       case pfCandidateCollection: {
0088         const auto& vetoes = e.get(vetoPFCandidatesSrc_);
0089         for (const auto& veto : vetoes) {
0090           if (veto.trackRef().isNull())
0091             continue;
0092           vetoed.emplace_back(veto.trackRef().id(), veto.trackRef().key());
0093         }
0094         break;
0095       }
0096     }  // switch
0097     std::sort(vetoed.begin(), vetoed.end());
0098   }
0099   const auto muonH = e.getHandle(muons_);
0100   const auto& muons = *muonH;
0101   elems.reserve(elems.size() + tracks->size());
0102   std::vector<bool> mask(tracks->size(), true);
0103   reco::MuonRef muonref;
0104
0105   // remove converted brems with bad pT resolution if requested
0106   // this reproduces the old behavior of PFBlockAlgo
0107   if (cleanBadConvBrems_) {
0108     auto itr = elems.begin();
0109     while (itr != elems.end()) {
0110       if ((*itr)->type() == reco::PFBlockElement::TRACK) {
0111         const reco::PFBlockElementTrack* trkel = static_cast<reco::PFBlockElementTrack*>(itr->get());
0112         const reco::ConversionRefVector& cRef = trkel->convRefs();
0113         const reco::PFDisplacedTrackerVertexRef& dvRef = trkel->displacedVertexRef(reco::PFBlockElement::T_FROM_DISP);
0114         const reco::VertexCompositeCandidateRef& v0Ref = trkel->V0Ref();
0115         // if there is no displaced vertex reference  and it is marked
0116         // as a conversion it's gotta be a converted brem
0117         if (trkel->trackType(reco::PFBlockElement::T_FROM_GAMMACONV) && cRef.empty() && dvRef.isNull() &&
0118             v0Ref.isNull()) {
0119           // if the Pt resolution is bad we kill this element
0120           if (!PFTrackAlgoTools::goodPtResolution(
0121                   trkel->trackRef(), DPtovPtCut_, NHitCut_, useIterTracking_, trackQuality_)) {
0122             itr = elems.erase(itr);
0123             continue;
0124           }
0125         }
0126       }
0127       ++itr;
0128     }  // loop on existing elements
0129   }
0130   // preprocess existing tracks in the element list and create a mask
0131   // so that we do not import tracks twice, tag muons we find
0132   // in this collection
0133   auto TKs_end = std::partition(
0134       elems.begin(), elems.end(), [](const ElementType& a) { return a->type() == reco::PFBlockElement::TRACK; });
0135   auto btk_elems = elems.begin();
0136   auto btrack = tracks->cbegin();
0137   auto etrack = tracks->cend();
0138   for (auto track = btrack; track != etrack; ++track) {
0139     auto tk_elem =
0140         std::find_if(btk_elems, TKs_end, [&](const ElementType& a) { return (a->trackRef() == track->trackRef()); });
0141     if (tk_elem != TKs_end) {
0142       mask[std::distance(tracks->cbegin(), track)] = false;
0143       // check and update if this track is a muon
0144       const int muId = PFMuonAlgo::muAssocToTrack((*tk_elem)->trackRef(), muons);
0145       if (muId != -1) {
0146         muonref = reco::MuonRef(muonH, muId);
0147         if (PFMuonAlgo::isLooseMuon(muonref) || PFMuonAlgo::isMuon(muonref)) {
0148           static_cast<reco::PFBlockElementTrack*>(tk_elem->get())->setMuonRef(muonref);
0149         }
0150       }
0151     }
0152   }
0153   // now we actually insert tracks, again tagging muons along the way
0154   reco::PFRecTrackRef pftrackref;
0155   reco::PFBlockElementTrack* trkElem = nullptr;
0156   for (auto track = btrack; track != etrack; ++track) {
0157     const unsigned idx = std::distance(btrack, track);
0158     // since we already set muon refs in the previously imported tracks,
0159     // here we can skip everything that is already imported
0160     if (!mask[idx])
0161       continue;
0162     muonref = reco::MuonRef();
0163     pftrackref = reco::PFRecTrackRef(tracks, idx);
0164     // Get the eventual muon associated to this track
0165     const int muId = PFMuonAlgo::muAssocToTrack(pftrackref->trackRef(), muons);
0166     bool thisIsAPotentialMuon = false;
0167     if (muId != -1) {
0168       muonref = reco::MuonRef(muonH, muId);
0169       thisIsAPotentialMuon =
0170           ((PFMuonAlgo::hasValidTrack(muonref, true, muonMaxDPtOPt_) && PFMuonAlgo::isLooseMuon(muonref)) ||
0171            (PFMuonAlgo::hasValidTrack(muonref, false, muonMaxDPtOPt_) && PFMuonAlgo::isMuon(muonref)));
0172     }
0173     if (thisIsAPotentialMuon || PFTrackAlgoTools::goodPtResolution(
0174                                     pftrackref->trackRef(), DPtovPtCut_, NHitCut_, useIterTracking_, trackQuality_)) {
0175       trkElem = new reco::PFBlockElementTrack(pftrackref);
0176       if (thisIsAPotentialMuon) {
0177         LogDebug("GeneralTracksImporter")
0178             << "Potential Muon P " << pftrackref->trackRef()->p() << " pt " << pftrackref->trackRef()->p() << std::endl;
0179       }
0180       if (muId != -1)
0181         trkElem->setMuonRef(muonref);
0182       //
0183       // if _vetoEndcap is false, add this trk automatically.
0184       // if _vetoEndcap is true, veto against the hgcal region tracks or charged PF candidates.
0185       // when simPF is used, we don't veto tracks with muonref even if they are in the hgcal region.
0186       //
0187       if (!vetoEndcap_)
0188         elems.emplace_back(trkElem);
0189       else {
0190         TrackProdIDKey trk = std::make_pair(pftrackref->trackRef().id(), pftrackref->trackRef().key());
0191         auto lower = std::lower_bound(vetoed.begin(), vetoed.end(), trk);
0192         bool inVetoList = (lower != vetoed.end() && *lower == trk);
0193         if (!inVetoList || (vetoMode_ == pfRecTrackCollection && muonref.isNonnull())) {
0194           elems.emplace_back(trkElem);
0195         } else
0196           delete trkElem;
0197       }
0198     }
0199   }
0200   elems.shrink_to_fit();
0201 }