RecoAlgos/interface/TrackFullCloneSelectorBase.h

0001 #ifndef RecoAlgos_TrackFullCloneSelectorBase_h
0002 #define RecoAlgos_TrackFullCloneSelectorBase_h
0003 /** \class TrackFullCloneSelectorBase
0004  *
0005  * selects a subset of a track collection, copying extra information on demand
0006  *
0007  * \author Giovanni Petrucciani
0008  *
0009  * \version $Revision: 1.3 $
0010  *
0011  * $Id: TrackFullCloneSelectorBase.h,v 1.3 2010/02/11 00:10:51 wmtan Exp $
0012  *
0013  */
0014
0015 #include <memory>
0016
0017 #include "FWCore/Framework/interface/ConsumesCollector.h"
0018 #include "FWCore/Framework/interface/Event.h"
0019 #include "FWCore/Framework/interface/stream/EDProducer.h"
0020 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0021 #include "FWCore/Utilities/interface/InputTag.h"
0022 #include <algorithm>
0023 #include <map>
0024 #include <memory>
0025 #include <utility>
0026 #include <vector>
0027
0028 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0029 #include "DataFormats/TrackReco/interface/Track.h"
0030 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0031 #include "TrackingTools/PatternTools/interface/Trajectory.h"
0032 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0033
0034 namespace reco {
0035   namespace modules {
0036
0037     template <typename Selector>
0038     class TrackFullCloneSelectorBase : public edm::stream::EDProducer<> {
0039     public:
0040       /// constructor
0041       explicit TrackFullCloneSelectorBase(const edm::ParameterSet& cfg)
0042           : hSrcTrackToken_(consumes<reco::TrackCollection>(cfg.template getParameter<edm::InputTag>("src"))),
0043             hTrajToken_(mayConsume<std::vector<Trajectory> >(cfg.template getParameter<edm::InputTag>("src"))),
0044             hTTAssToken_(mayConsume<TrajTrackAssociationCollection>(cfg.template getParameter<edm::InputTag>("src"))),
0045             copyExtras_(cfg.template getUntrackedParameter<bool>("copyExtras", false)),
0046             copyTrajectories_(cfg.template getUntrackedParameter<bool>("copyTrajectories", false)),
0047             selector_(cfg, consumesCollector()) {
0048         std::string alias(cfg.getParameter<std::string>("@module_label"));
0049         produces<reco::TrackCollection>().setBranchAlias(alias + "Tracks");
0050         if (copyExtras_) {
0051           produces<reco::TrackExtraCollection>().setBranchAlias(alias + "TrackExtras");
0052           produces<TrackingRecHitCollection>().setBranchAlias(alias + "RecHits");
0053           if (copyTrajectories_) {
0054             produces<std::vector<Trajectory> >().setBranchAlias(alias + "Trajectories");
0055             produces<TrajTrackAssociationCollection>().setBranchAlias(alias + "TrajectoryTrackAssociations");
0056           }
0057         }
0058       }
0059       /// destructor
0060       ~TrackFullCloneSelectorBase() override {}
0061
0062     private:
0063       /// process one event
0064       void produce(edm::Event& evt, const edm::EventSetup& es) override {
0065         edm::Handle<reco::TrackCollection> hSrcTrack;
0066         evt.getByToken(hSrcTrackToken_, hSrcTrack);
0067
0068         selTracks_ = std::make_unique<reco::TrackCollection>();
0069         if (copyExtras_) {
0070           selTrackExtras_ = std::make_unique<reco::TrackExtraCollection>();
0071           selHits_ = std::make_unique<TrackingRecHitCollection>();
0072         }
0073
0074         TrackRefProd rTracks = evt.template getRefBeforePut<TrackCollection>();
0075
0076         TrackingRecHitRefProd rHits;
0077         TrackExtraRefProd rTrackExtras;
0078         if (copyExtras_) {
0079           rHits = evt.template getRefBeforePut<TrackingRecHitCollection>();
0080           rTrackExtras = evt.template getRefBeforePut<TrackExtraCollection>();
0081         }
0082
0083         typedef reco::TrackRef::key_type TrackRefKey;
0084         std::map<TrackRefKey, reco::TrackRef> goodTracks;
0085         TrackRefKey current = 0;
0086
0087         selector_.init(evt, es);
0088         auto tkBegin = hSrcTrack->begin();
0089         for (reco::TrackCollection::const_iterator it = tkBegin, ed = hSrcTrack->end(); it != ed; ++it, ++current) {
0090           const reco::Track& trk = *it;
0091           const reco::TrackRef tkref(hSrcTrack, std::distance(tkBegin, it));
0092           if (!selector_(tkref))
0093             continue;
0094
0095           selTracks_->push_back(Track(trk));  // clone and store
0096           if (!copyExtras_)
0097             continue;
0098
0099           // TrackExtras
0100           selTrackExtras_->push_back(TrackExtra(trk.outerPosition(),
0101                                                 trk.outerMomentum(),
0102                                                 trk.outerOk(),
0103                                                 trk.innerPosition(),
0104                                                 trk.innerMomentum(),
0105                                                 trk.innerOk(),
0106                                                 trk.outerStateCovariance(),
0107                                                 trk.outerDetId(),
0108                                                 trk.innerStateCovariance(),
0109                                                 trk.innerDetId(),
0110                                                 trk.seedDirection()));
0111           selTracks_->back().setExtra(TrackExtraRef(rTrackExtras, selTrackExtras_->size() - 1));
0112           TrackExtra& tx = selTrackExtras_->back();
0113           // TrackingRecHits
0114           auto const firstHitIndex = selHits_->size();
0115           for (auto const& hit : trk.recHits())
0116             selHits_->push_back(hit->clone());
0117           tx.setHits(rHits, firstHitIndex, selHits_->size() - firstHitIndex);
0118           tx.setTrajParams(trk.extra()->trajParams(), trk.extra()->chi2sX5());
0119           assert(tx.trajParams().size() == tx.recHitsSize());
0120           if (copyTrajectories_) {
0121             goodTracks[current] = reco::TrackRef(rTracks, selTracks_->size() - 1);
0122           }
0123         }
0124         if (copyTrajectories_) {
0125           edm::Handle<std::vector<Trajectory> > hTraj;
0126           edm::Handle<TrajTrackAssociationCollection> hTTAss;
0127           evt.getByToken(hTTAssToken_, hTTAss);
0128           evt.getByToken(hTrajToken_, hTraj);
0129           edm::RefProd<std::vector<Trajectory> > TrajRefProd = evt.template getRefBeforePut<std::vector<Trajectory> >();
0130           selTrajs_ = std::make_unique<std::vector<Trajectory> >();
0131           selTTAss_ = std::make_unique<TrajTrackAssociationCollection>();
0132           for (size_t i = 0, n = hTraj->size(); i < n; ++i) {
0133             edm::Ref<std::vector<Trajectory> > trajRef(hTraj, i);
0134             TrajTrackAssociationCollection::const_iterator match = hTTAss->find(trajRef);
0135             if (match != hTTAss->end()) {
0136               const edm::Ref<reco::TrackCollection>& trkRef = match->val;
0137               TrackRefKey oldKey = trkRef.key();
0138               std::map<TrackRefKey, reco::TrackRef>::iterator getref = goodTracks.find(oldKey);
0139               if (getref != goodTracks.end()) {
0140                 // do the clone
0141                 selTrajs_->push_back(Trajectory(*trajRef));
0142                 selTTAss_->insert(edm::Ref<std::vector<Trajectory> >(TrajRefProd, selTrajs_->size() - 1),
0143                                   getref->second);
0144               }
0145             }
0146           }
0147         }
0148
0149         evt.put(std::move(selTracks_));
0150         if (copyExtras_) {
0151           evt.put(std::move(selTrackExtras_));
0152           evt.put(std::move(selHits_));
0153           if (copyTrajectories_) {
0154             evt.put(std::move(selTrajs_));
0155             evt.put(std::move(selTTAss_));
0156           }
0157         }
0158       }
0159       /// source collection label
0160       edm::EDGetTokenT<reco::TrackCollection> hSrcTrackToken_;
0161       edm::EDGetTokenT<std::vector<Trajectory> > hTrajToken_;
0162       edm::EDGetTokenT<TrajTrackAssociationCollection> hTTAssToken_;
0163       /// copy only the tracks, not extras and rechits (for AOD)
0164       bool copyExtras_;
0165       /// copy also trajectories and trajectory->track associations
0166       bool copyTrajectories_;
0167       /// filter event
0168       Selector selector_;
0169       // some space
0170       std::unique_ptr<reco::TrackCollection> selTracks_;
0171       std::unique_ptr<reco::TrackExtraCollection> selTrackExtras_;
0172       std::unique_ptr<TrackingRecHitCollection> selHits_;
0173       std::unique_ptr<std::vector<Trajectory> > selTrajs_;
0174       std::unique_ptr<TrajTrackAssociationCollection> selTTAss_;
0175     };
0176
0177   }  // namespace modules
0178 }  // namespace reco
0179 #endif