PixelLowPtUtilities/plugins/TrackListCombiner.cc

0001 #include "TrackListCombiner.h"
0002
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/Framework/interface/EventSetup.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0006
0007 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0008
0009 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0010 #include "DataFormats/TrackReco/interface/Track.h"
0011 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0012
0013 using namespace std;
0014
0015 /*****************************************************************************/
0016 TrackListCombiner::TrackListCombiner(const edm::ParameterSet& ps) {
0017   for (const std::string& prod : ps.getParameter<vector<string>>("trackProducers")) {
0018     trackProducers.emplace_back(consumes<vector<Trajectory>>(prod), consumes<TrajTrackAssociationCollection>(prod));
0019   }
0020
0021   produces<reco::TrackCollection>();
0022   produces<reco::TrackExtraCollection>();
0023   produces<TrackingRecHitCollection>();
0024   produces<vector<Trajectory>>();
0025   produces<TrajTrackAssociationCollection>();
0026 }
0027
0028 /*****************************************************************************/
0029 TrackListCombiner::~TrackListCombiner() {}
0030
0031 /*****************************************************************************/
0032 void TrackListCombiner::produce(edm::StreamID, edm::Event& ev, const edm::EventSetup& es) const {
0033   auto recoTracks = std::make_unique<reco::TrackCollection>();
0034   auto recoTrackExtras = std::make_unique<reco::TrackExtraCollection>();
0035   auto recoHits = std::make_unique<TrackingRecHitCollection>();
0036   auto recoTrajectories = std::make_unique<vector<Trajectory>>();
0037   auto recoTrajTrackMap = std::make_unique<TrajTrackAssociationCollection>();
0038
0039   LogTrace("MinBiasTracking") << "[TrackListCombiner]";
0040
0041   // Go through all track producers
0042   int i = 1;
0043   for (auto trackProducer = trackProducers.begin(); trackProducer != trackProducers.end(); trackProducer++, i++) {
0044     reco::TrackBase::TrackAlgorithm algo;
0045     switch (i) {
0046       case 1:
0047         algo = reco::TrackBase::lowPtTripletStep;
0048         break;
0049       case 2:
0050         algo = reco::TrackBase::pixelPairStep;
0051         break;
0052       case 3:
0053         algo = reco::TrackBase::detachedTripletStep;
0054         break;
0055       default:
0056         algo = reco::TrackBase::undefAlgorithm;
0057     }
0058
0059     edm::Handle<vector<Trajectory>> theTrajectoryCollection;
0060     edm::Handle<TrajTrackAssociationCollection> theAssoMap;
0061
0062     ev.getByToken(trackProducer->trajectory, theTrajectoryCollection);
0063     ev.getByToken(trackProducer->assoMap, theAssoMap);
0064
0065 #ifdef EDM_ML_DEBUG
0066     edm::EDConsumerBase::Labels labels;
0067     labelsForToken(trackProducer->trajectory, labels);
0068
0069     LogTrace("MinBiasTracking") << " [TrackListCombiner] " << labels.module << " : " << theAssoMap->size();
0070 #endif
0071
0072     // The track collection iterators
0073     TrajTrackAssociationCollection::const_iterator anAssociation;
0074     TrajTrackAssociationCollection::const_iterator lastAssociation;
0075     anAssociation = theAssoMap->begin();
0076     lastAssociation = theAssoMap->end();
0077
0078     // Build the map of correspondance between reco tracks and sim tracks
0079     for (; anAssociation != lastAssociation; ++anAssociation) {
0080       edm::Ref<vector<Trajectory>> aTrajectoryRef = anAssociation->key;
0081       reco::TrackRef aTrackRef = anAssociation->val;
0082
0083       // A copy of the track
0084       reco::Track aRecoTrack(*aTrackRef);
0085
0086       // Set algorithm
0087       aRecoTrack.setAlgorithm(algo);
0088
0089       recoTracks->push_back(aRecoTrack);
0090
0091       // A copy of the hits
0092       unsigned nh = aRecoTrack.recHitsSize();
0093       for (unsigned ih = 0; ih < nh; ++ih) {
0094         TrackingRecHit* hit = aRecoTrack.recHit(ih)->clone();
0095         recoHits->push_back(hit);
0096       }
0097
0098       // A copy of the trajectories
0099       recoTrajectories->push_back(*aTrajectoryRef);
0100     }
0101   }
0102
0103   LogTrace("MinBiasTracking") << " [TrackListCombiner] allTracks : " << recoTracks->size() << "|"
0104                               << recoTrajectories->size();
0105
0106   // Save the tracking recHits
0107   edm::OrphanHandle<TrackingRecHitCollection> theRecoHits = ev.put(std::move(recoHits));
0108
0109   edm::RefProd<TrackingRecHitCollection> theRecoHitsProd(theRecoHits);
0110   // Create the track extras and add the references to the rechits
0111   unsigned hits = 0;
0112   unsigned nTracks = recoTracks->size();
0113   recoTrackExtras->reserve(nTracks);  // To save some time at push_back
0114   for (unsigned index = 0; index < nTracks; ++index) {
0115     reco::Track& aTrack = recoTracks->at(index);
0116     reco::TrackExtra aTrackExtra(aTrack.outerPosition(),
0117                                  aTrack.outerMomentum(),
0118                                  aTrack.outerOk(),
0119                                  aTrack.innerPosition(),
0120                                  aTrack.innerMomentum(),
0121                                  aTrack.innerOk(),
0122                                  aTrack.outerStateCovariance(),
0123                                  aTrack.outerDetId(),
0124                                  aTrack.innerStateCovariance(),
0125                                  aTrack.innerDetId(),
0126                                  aTrack.seedDirection(),
0127                                  aTrack.seedRef());
0128
0129     unsigned nHits = aTrack.recHitsSize();
0130     aTrackExtra.setHits(theRecoHitsProd, hits, nHits);
0131     hits += nHits;
0132
0133     recoTrackExtras->push_back(aTrackExtra);
0134   }
0135
0136   // Save the track extras
0137   edm::OrphanHandle<reco::TrackExtraCollection> theRecoTrackExtras = ev.put(std::move(recoTrackExtras));
0138
0139   // Add the reference to the track extra in the tracks
0140   for (unsigned index = 0; index < nTracks; ++index) {
0141     const reco::TrackExtraRef theTrackExtraRef(theRecoTrackExtras, index);
0142     (recoTracks->at(index)).setExtra(theTrackExtraRef);
0143   }
0144
0145   // Save the tracks
0146   edm::OrphanHandle<reco::TrackCollection> theRecoTracks = ev.put(std::move(recoTracks));
0147
0148   // Save the trajectories
0149   edm::OrphanHandle<vector<Trajectory>> theRecoTrajectories = ev.put(std::move(recoTrajectories));
0150
0151   // Create and set the trajectory/track association map
0152   for (unsigned index = 0; index < nTracks; ++index) {
0153     edm::Ref<vector<Trajectory>> trajRef(theRecoTrajectories, index);
0154     edm::Ref<reco::TrackCollection> tkRef(theRecoTracks, index);
0155     recoTrajTrackMap->insert(trajRef, tkRef);
0156   }
0157
0158   // Save the association map
0159   ev.put(std::move(recoTrajTrackMap));
0160 }