TrackProducer/plugins/QualityFilter.cc

0001 #include <memory>
0002
0003 // user include files
0004 #include "FWCore/Framework/interface/Frameworkfwd.h"
0005 #include "FWCore/Framework/interface/stream/EDProducer.h"
0006 #include "DataFormats/TrackReco/interface/Track.h"
0007 #include "FWCore/Framework/interface/Event.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0010
0011 class dso_hidden QualityFilter final : public edm::stream::EDProducer<> {
0012 public:
0013   explicit QualityFilter(const edm::ParameterSet&);
0014   ~QualityFilter() override;
0015
0016 private:
0017   void produce(edm::Event&, const edm::EventSetup&) override;
0018   virtual void endJob();
0019
0020   // ----------member data ---------------------------
0021 private:
0022   edm::EDGetTokenT<std::vector<Trajectory> > trajTag;
0023   edm::EDGetTokenT<TrajTrackAssociationCollection> tassTag;
0024   reco::TrackBase::TrackQuality trackQuality_;
0025   bool copyExtras_;
0026 };
0027
0028 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0029 #include "DataFormats/TrackReco/interface/TrackExtraFwd.h"
0030 #include "FWCore/Utilities/interface/InputTag.h"
0031 #include "TrackingTools/TransientTrackingRecHit/interface/TransientTrackingRecHit.h"
0032 #include "TrackingTools/PatternTools/interface/Trajectory.h"
0033 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0034
0035 //
0036 // class decleration
0037 //
0038
0039 using namespace edm;
0040 using namespace reco;
0041 using namespace std;
0042
0043 QualityFilter::QualityFilter(const edm::ParameterSet& iConfig) {
0044   copyExtras_ = iConfig.getUntrackedParameter<bool>("copyExtras", false);
0045
0046   produces<reco::TrackCollection>();
0047   if (copyExtras_) {
0048     produces<TrackingRecHitCollection>();
0049     produces<reco::TrackExtraCollection>();
0050   }
0051   produces<std::vector<Trajectory> >();
0052   produces<TrajTrackAssociationCollection>();
0053
0054   trajTag = consumes<std::vector<Trajectory> >(iConfig.getParameter<edm::InputTag>("recTracks"));
0055   tassTag = consumes<TrajTrackAssociationCollection>(iConfig.getParameter<edm::InputTag>("recTracks"));
0056   trackQuality_ = TrackBase::qualityByName(iConfig.getParameter<std::string>("TrackQuality"));
0057 }
0058
0059 QualityFilter::~QualityFilter() {
0060   // do anything here that needs to be done at desctruction time
0061   // (e.g. close files, deallocate resources etc.)
0062 }
0063
0064 //
0065 // member functions
0066 //
0067
0068 // ------------ method called to produce the data  ------------
0069 void QualityFilter::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0070   unique_ptr<TrackCollection> selTracks(new TrackCollection);
0071   unique_ptr<TrackingRecHitCollection> selHits(copyExtras_ ? new TrackingRecHitCollection() : nullptr);
0072   unique_ptr<TrackExtraCollection> selTrackExtras(copyExtras_ ? new TrackExtraCollection() : nullptr);
0073   unique_ptr<vector<Trajectory> > outputTJ(new vector<Trajectory>);
0074   unique_ptr<TrajTrackAssociationCollection> trajTrackMap(new TrajTrackAssociationCollection());
0075
0076   TrackExtraRefProd rTrackExtras;
0077   TrackingRecHitRefProd rHits;
0078   if (copyExtras_) {
0079     rTrackExtras = iEvent.getRefBeforePut<TrackExtraCollection>();
0080     rHits = iEvent.getRefBeforePut<TrackingRecHitCollection>();
0081   }
0082
0083   Handle<std::vector<Trajectory> > TrajectoryCollection;
0084   Handle<TrajTrackAssociationCollection> assoMap;
0085
0086   iEvent.getByToken(trajTag, TrajectoryCollection);
0087   iEvent.getByToken(tassTag, assoMap);
0088
0089   TrajTrackAssociationCollection::const_iterator it = assoMap->begin();
0090   TrajTrackAssociationCollection::const_iterator lastAssoc = assoMap->end();
0091   for (; it != lastAssoc; ++it) {
0092     const Ref<vector<Trajectory> > traj = it->key;
0093     const reco::TrackRef itc = it->val;
0094     bool goodTk = (itc->quality(trackQuality_));
0095
0096     if (goodTk) {
0097       auto const& track = (*itc);
0098       //tracks and trajectories
0099       selTracks->push_back(track);
0100       outputTJ->push_back(*traj);
0101       if (copyExtras_) {
0102         //TRACKING HITS
0103         trackingRecHit_iterator irhit = (*itc).recHitsBegin();
0104         trackingRecHit_iterator lasthit = (*itc).recHitsEnd();
0105         for (; irhit != lasthit; ++irhit) {
0106           selHits->push_back((*irhit)->clone());
0107         }
0108       }
0109     }
0110   }
0111
0112   unsigned nTracks = selTracks->size();
0113   if (copyExtras_) {
0114     //PUT TRACKING HITS IN THE EVENT
0115     OrphanHandle<TrackingRecHitCollection> theRecoHits = iEvent.put(std::move(selHits));
0116     edm::RefProd<TrackingRecHitCollection> theRecoHitsProd(theRecoHits);
0117
0118     //PUT TRACK EXTRA IN THE EVENT
0119     selTrackExtras->reserve(nTracks);
0120     unsigned hits = 0;
0121
0122     for (unsigned index = 0; index < nTracks; ++index) {
0123       auto const& aTrack = (*selTracks)[index];
0124       selTrackExtras->emplace_back(aTrack.outerPosition(),
0125                                    aTrack.outerMomentum(),
0126                                    aTrack.outerOk(),
0127                                    aTrack.innerPosition(),
0128                                    aTrack.innerMomentum(),
0129                                    aTrack.innerOk(),
0130                                    aTrack.outerStateCovariance(),
0131                                    aTrack.outerDetId(),
0132                                    aTrack.innerStateCovariance(),
0133                                    aTrack.innerDetId(),
0134                                    aTrack.seedDirection(),
0135                                    aTrack.seedRef());
0136
0137       auto& aTrackExtra = selTrackExtras->back();
0138       //unsigned nHits = aTrack.numberOfValidHits();
0139       auto nHits = aTrack.recHitsSize();
0140       aTrackExtra.setHits(theRecoHitsProd, hits, nHits);
0141       hits += nHits;
0142       selTrackExtras->push_back(aTrackExtra);
0143     }
0144
0145     //CORRECT REF TO TRACK
0146     OrphanHandle<TrackExtraCollection> theRecoTrackExtras = iEvent.put(std::move(selTrackExtras));
0147     for (unsigned index = 0; index < nTracks; ++index) {
0148       const reco::TrackExtraRef theTrackExtraRef(theRecoTrackExtras, index);
0149       (*selTracks)[index].setExtra(theTrackExtraRef);
0150     }
0151   }  // END IF COPY EXTRAS
0152
0153   //TRACKS AND TRAJECTORIES
0154   OrphanHandle<TrackCollection> theRecoTracks = iEvent.put(std::move(selTracks));
0155   OrphanHandle<vector<Trajectory> > theRecoTrajectories = iEvent.put(std::move(outputTJ));
0156
0157   //TRACKS<->TRAJECTORIES MAP
0158   nTracks = theRecoTracks->size();
0159   for (unsigned index = 0; index < nTracks; ++index) {
0160     Ref<vector<Trajectory> > trajRef(theRecoTrajectories, index);
0161     Ref<TrackCollection> tkRef(theRecoTracks, index);
0162     trajTrackMap->insert(trajRef, tkRef);
0163   }
0164   //MAP IN THE EVENT
0165   iEvent.put(std::move(trajTrackMap));
0166 }
0167
0168 // ------------ method called once each job just after ending the event loop  ------------
0169 void QualityFilter::endJob() {}
0170
0171 DEFINE_FWK_MODULE(QualityFilter);