MuonIdentification/plugins/MuonReducedTrackExtraProducer.cc

0001 #include "RecoMuon/MuonIdentification/plugins/MuonReducedTrackExtraProducer.h"
0002 #include "DataFormats/MuonReco/interface/Muon.h"
0003 #include "DataFormats/TrackerRecHit2D/interface/TrackerSingleRecHit.h"
0004 #include "DataFormats/TrackerRecHit2D/interface/SiStripMatchedRecHit2D.h"
0005
0006 MuonReducedTrackExtraProducer::MuonReducedTrackExtraProducer(const edm::ParameterSet& pset)
0007     : muonToken_(consumes<edm::View<reco::Muon>>(pset.getParameter<edm::InputTag>("muonTag"))),
0008       outputClusters_(pset.getParameter<bool>("outputClusters")),
0009       selector_(pset.getParameter<std::string>("cut")),
0010       trackExtraOutToken_(produces<reco::TrackExtraCollection>()),
0011       trackingRecHitsOutToken_(produces<TrackingRecHitCollection>()),
0012       associationOutToken_(produces<edm::Association<reco::TrackExtraCollection>>()) {
0013   std::vector<edm::InputTag> trackExtraTags = pset.getParameter<std::vector<edm::InputTag>>("trackExtraTags");
0014   for (edm::InputTag const& tag : trackExtraTags) {
0015     trackExtraTokens_.push_back(consumes<reco::TrackExtraCollection>(tag));
0016   }
0017
0018   std::vector<edm::InputTag> trackExtraAssocTags = pset.getParameter<std::vector<edm::InputTag>>("trackExtraAssocs");
0019   for (edm::InputTag const& tag : trackExtraAssocTags) {
0020     trackExtraAssocs_.push_back(consumes<edm::Association<reco::TrackExtraCollection>>(tag));
0021   }
0022
0023   if (outputClusters_) {
0024     pixelClusterToken_ =
0025         consumes<edmNew::DetSetVector<SiPixelCluster>>(pset.getParameter<edm::InputTag>("pixelClusterTag"));
0026     stripClusterToken_ =
0027         consumes<edmNew::DetSetVector<SiStripCluster>>(pset.getParameter<edm::InputTag>("stripClusterTag"));
0028
0029     pixelClusterOutToken_ = produces<edmNew::DetSetVector<SiPixelCluster>>();
0030     stripClusterOutToken_ = produces<edmNew::DetSetVector<SiStripCluster>>();
0031   }
0032 }
0033
0034 void MuonReducedTrackExtraProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0035   edm::ParameterSetDescription desc;
0036   desc.setComment(
0037       "Produces reduced set of TrackExtras and corresponding TrackingRecHits and (optionally) Pixe/Strip clusters "
0038       "associated to a muon track.");
0039   desc.add<edm::InputTag>("muonTag", edm::InputTag("muons"));
0040   desc.add<std::vector<edm::InputTag>>("trackExtraTags",
0041                                        {edm::InputTag("generalTracks"),
0042                                         edm::InputTag("globalMuons"),
0043                                         edm::InputTag("tevMuons", "firstHit"),
0044                                         edm::InputTag("tevMuons", "picky"),
0045                                         edm::InputTag("tevMuons", "dyt")});
0046   desc.add<std::vector<edm::InputTag>>("trackExtraAssocs", {});
0047   desc.add<edm::InputTag>("pixelClusterTag", edm::InputTag("siPixelClusters"));
0048   desc.add<edm::InputTag>("stripClusterTag", edm::InputTag("siStripClusters"));
0049   desc.add<std::string>("cut", "pt > 3.0");
0050   desc.add<bool>("outputClusters", true);
0051   descriptions.add("muonReducedTrackExtras", desc);
0052 }
0053
0054 void MuonReducedTrackExtraProducer::produce(edm::Event& event, const edm::EventSetup& eventSetup) {
0055   auto muons = event.getHandle(muonToken_);
0056
0057   std::vector<edm::Handle<reco::TrackExtraCollection>> trackExtrasV(trackExtraTokens_.size());
0058   for (unsigned int i = 0; i < trackExtraTokens_.size(); ++i) {
0059     event.getByToken(trackExtraTokens_[i], trackExtrasV[i]);
0060   }
0061
0062   std::vector<edm::Handle<edm::Association<reco::TrackExtraCollection>>> trackExtraAssocs(trackExtraAssocs_.size());
0063   for (unsigned int i = 0; i < trackExtraAssocs_.size(); ++i) {
0064     event.getByToken(trackExtraAssocs_[i], trackExtraAssocs[i]);
0065   }
0066
0067   std::map<edm::ProductID, std::vector<bool>> idxstokeep;
0068   for (auto const& trackExtras : trackExtrasV) {
0069     idxstokeep[trackExtras.id()].resize(trackExtras->size(), false);
0070   }
0071
0072   //loop over muons and mark track extras to keep
0073   for (auto const& muon : *muons) {
0074     if (!selector_(muon)) {
0075       continue;
0076     }
0077     reco::TrackExtraRef trackExtra = muon.bestTrack()->extra();
0078     // check recursively through association maps if present
0079     for (auto const& assoc : trackExtraAssocs) {
0080       if (!assoc->contains(trackExtra.id())) {
0081         continue;
0082       }
0083       reco::TrackExtraRef const& trackExtraOut = (*assoc)[trackExtra];
0084       if (trackExtraOut.isNonnull()) {
0085         trackExtra = trackExtraOut;
0086       }
0087     }
0088     auto idxs = idxstokeep.find(trackExtra.id());
0089     if (idxs != idxstokeep.end()) {
0090       idxs->second[trackExtra.key()] = true;
0091     }
0092   }
0093
0094   //output collections for TrackExtras and TrackingRecHits and the association map
0095   reco::TrackExtraCollection trackExtrasOut;
0096   TrackingRecHitCollection trackingRecHitsOut;
0097   edm::Association<reco::TrackExtraCollection> association;
0098   edm::Association<reco::TrackExtraCollection>::Filler assocfiller(association);
0099
0100   //refprod for the output TrackExtraCollection
0101   edm::RefProd<reco::TrackExtraCollection> trackExtraRefProd = event.getRefBeforePut(trackExtraOutToken_);
0102   //refprod for the output TrackingRecHitCollection
0103   edm::RefProd<TrackingRecHitCollection> hitRefProd = event.getRefBeforePut(trackingRecHitsOutToken_);
0104
0105   association.setRef(trackExtraRefProd);
0106
0107   edm::Handle<edmNew::DetSetVector<SiPixelCluster>> pixelClusters;
0108   edm::Handle<edmNew::DetSetVector<SiStripCluster>> stripClusters;
0109
0110   //indexes of pixel/strip clusters to keep
0111   std::vector<bool> pixelClustersToKeep;
0112   std::vector<bool> stripClustersToKeep;
0113
0114   if (outputClusters_) {
0115     event.getByToken(pixelClusterToken_, pixelClusters);
0116     event.getByToken(stripClusterToken_, stripClusters);
0117
0118     pixelClustersToKeep.resize(pixelClusters->dataSize(), false);
0119     stripClustersToKeep.resize(stripClusters->dataSize(), false);
0120   }
0121
0122   using SiPixelClusterRef = edm::Ref<edmNew::DetSetVector<SiPixelCluster>, SiPixelCluster>;
0123   using SiStripClusterRef = edm::Ref<edmNew::DetSetVector<SiStripCluster>, SiStripCluster>;
0124
0125   //loop over track extras and fill output together with TrackingRechits
0126   //as well as marking pixel and strip clusters to keep
0127   for (auto const& trackExtras : trackExtrasV) {
0128     const std::vector<bool>& idxs = idxstokeep.at(trackExtras.id());
0129     //indices for association (-1 for null association)
0130     std::vector<int> associdxs(trackExtras->size(), -1);
0131     for (unsigned int i = 0; i < trackExtras->size(); ++i) {
0132       if (!idxs[i]) {
0133         continue;
0134       }
0135       const reco::TrackExtra& trackExtra = (*trackExtras)[i];
0136
0137       //fill association idx
0138       associdxs[i] = trackExtrasOut.size();
0139
0140       //fill TrackExtra
0141       trackExtrasOut.emplace_back(trackExtra.outerPosition(),
0142                                   trackExtra.outerMomentum(),
0143                                   trackExtra.outerOk(),
0144                                   trackExtra.innerPosition(),
0145                                   trackExtra.innerMomentum(),
0146                                   trackExtra.innerOk(),
0147                                   trackExtra.outerStateCovariance(),
0148                                   trackExtra.outerDetId(),
0149                                   trackExtra.innerStateCovariance(),
0150                                   trackExtra.innerDetId(),
0151                                   trackExtra.seedDirection(),
0152                                   trackExtra.seedRef());
0153
0154       //rekey refs to TrackingRecHits
0155       reco::TrackExtra& trackExtraOut = trackExtrasOut.back();
0156       trackExtraOut.setHits(hitRefProd, trackingRecHitsOut.size(), trackExtra.recHitsSize());
0157       for (auto const& hit : trackExtra.recHits()) {
0158         if (outputClusters_) {
0159           //mark clusters to keep
0160           TrackerSingleRecHit const* singleHit = dynamic_cast<TrackerSingleRecHit const*>(&*hit);
0161           if (singleHit != nullptr) {
0162             SiPixelClusterRef const& pixelRef = singleHit->cluster_pixel();
0163             if (pixelRef.isNonnull() && pixelRef.id() == pixelClusters.id()) {
0164               pixelClustersToKeep[pixelRef.key()] = true;
0165             }
0166             SiStripClusterRef const& stripRef = singleHit->cluster_strip();
0167             if (stripRef.isNonnull() && stripRef.id() == stripClusters.id()) {
0168               stripClustersToKeep[stripRef.key()] = true;
0169             }
0170             SiStripMatchedRecHit2D const* matched2DHit = dynamic_cast<SiStripMatchedRecHit2D const*>(singleHit);
0171             if (matched2DHit != nullptr) {
0172               SiStripClusterRef const& monoRef = matched2DHit->monoClusterRef().cluster_strip();
0173               SiStripClusterRef const& stereoRef = matched2DHit->stereoClusterRef().cluster_strip();
0174               if (monoRef.isNonnull() && monoRef.id() == stripClusters.id()) {
0175                 stripClustersToKeep[monoRef.key()] = true;
0176               }
0177               if (stereoRef.isNonnull() && stereoRef.id() == stripClusters.id()) {
0178                 stripClustersToKeep[stereoRef.key()] = true;
0179               }
0180             }
0181           }
0182         }
0183         //fill output hit
0184         trackingRecHitsOut.push_back(hit->clone());
0185       }
0186     }
0187     assocfiller.insert(trackExtras, associdxs.begin(), associdxs.end());
0188   }
0189
0190   assocfiller.fill();
0191
0192   if (outputClusters_) {
0193     //output collections for clusters
0194     edmNew::DetSetVector<SiPixelCluster> pixelClustersOut;
0195     edmNew::DetSetVector<SiStripCluster> stripClustersOut;
0196     //mapping of indices from input to output collections
0197     std::unordered_map<unsigned int, unsigned int> pixelClusterIdxMap;
0198     std::unordered_map<unsigned int, unsigned int> stripClusterIdxMap;
0199
0200     //fill output clusters
0201     //this indexes the internal data array of the DetSetVector
0202     unsigned int iIndex = 0;
0203     //loop over edmNew:::DetSet<T>
0204     for (auto setIter = pixelClusters->begin(), setIterEnd = pixelClusters->end(); setIter != setIterEnd; ++setIter) {
0205       //fill items from this DetSet
0206       typename edmNew::DetSetVector<SiPixelCluster>::FastFiller ff(pixelClustersOut, setIter->detId());
0207       for (auto iter = setIter->begin(), iterEnd = setIter->end(); iter != iterEnd; ++iter, ++iIndex) {
0208         if (pixelClustersToKeep[iIndex]) {
0209           ff.push_back(*iter);
0210           const unsigned int oIndex = pixelClusterIdxMap.size();
0211           pixelClusterIdxMap[iIndex] = oIndex;
0212         }
0213       }
0214     }
0215
0216     iIndex = 0;
0217     for (auto setIter = stripClusters->begin(), setIterEnd = stripClusters->end(); setIter != setIterEnd; ++setIter) {
0218       //fill items from this DetSet
0219       typename edmNew::DetSetVector<SiStripCluster>::FastFiller ff(stripClustersOut, setIter->detId());
0220       for (auto iter = setIter->begin(), iterEnd = setIter->end(); iter != iterEnd; ++iter, ++iIndex) {
0221         if (stripClustersToKeep[iIndex]) {
0222           ff.push_back(*iter);
0223           const unsigned int oIndex = stripClusterIdxMap.size();
0224           stripClusterIdxMap[iIndex] = oIndex;
0225         }
0226       }
0227     }
0228
0229     edm::OrphanHandle<edmNew::DetSetVector<SiPixelCluster>> pixelClustersOutH =
0230         event.emplace(pixelClusterOutToken_, std::move(pixelClustersOut));
0231     edm::OrphanHandle<edmNew::DetSetVector<SiStripCluster>> stripClustersOutH =
0232         event.emplace(stripClusterOutToken_, std::move(stripClustersOut));
0233
0234     //rekey cluster references in output hit collection
0235     for (auto& hit : trackingRecHitsOut) {
0236       TrackerSingleRecHit* singleHit = dynamic_cast<TrackerSingleRecHit*>(&hit);
0237       if (singleHit != nullptr) {
0238         SiPixelClusterRef const& pixelRef = singleHit->cluster_pixel();
0239         if (pixelRef.isNonnull() && pixelRef.id() == pixelClusters.id()) {
0240           SiPixelClusterRef const pixelRefOut(pixelClustersOutH, pixelClusterIdxMap.at(pixelRef.key()));
0241           singleHit->setClusterPixelRef(pixelRefOut);
0242         }
0243         SiStripClusterRef const& stripRef = singleHit->cluster_strip();
0244         if (stripRef.isNonnull() && stripRef.id() == stripClusters.id()) {
0245           SiStripClusterRef const stripRefOut(stripClustersOutH, stripClusterIdxMap.at(stripRef.key()));
0246           singleHit->setClusterStripRef(stripRefOut);
0247         }
0248         SiStripMatchedRecHit2D* matched2DHit = dynamic_cast<SiStripMatchedRecHit2D*>(singleHit);
0249         if (matched2DHit != nullptr) {
0250           SiStripClusterRef const& monoRef = matched2DHit->monoClusterRef().cluster_strip();
0251           SiStripClusterRef const& stereoRef = matched2DHit->stereoClusterRef().cluster_strip();
0252           if (monoRef.isNonnull() && monoRef.id() == stripClusters.id()) {
0253             SiStripClusterRef const monoRefOut(stripClustersOutH, stripClusterIdxMap.at(monoRef.key()));
0254             matched2DHit->monoClusterRef() = OmniClusterRef(monoRefOut);
0255           }
0256           if (stereoRef.isNonnull() && stereoRef.id() == stripClusters.id()) {
0257             SiStripClusterRef const stereoRefOut(stripClustersOutH, stripClusterIdxMap.at(stereoRef.key()));
0258             matched2DHit->stereoClusterRef() = OmniClusterRef(stereoRefOut);
0259           }
0260         }
0261       }
0262     }
0263   }
0264
0265   event.emplace(trackExtraOutToken_, std::move(trackExtrasOut));
0266   event.emplace(trackingRecHitsOutToken_, std::move(trackingRecHitsOut));
0267   event.emplace(associationOutToken_, std::move(association));
0268 }