Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2025-01-08 03:35:56

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 <algorithm>
0016 #include <map>
0017 #include <memory>
0018 #include <utility>
0019 #include <vector>
0020 
0021 #include "DataFormats/TrackReco/interface/Track.h"
0022 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0023 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0024 #include "FWCore/Framework/interface/ConsumesCollector.h"
0025 #include "FWCore/Framework/interface/Event.h"
0026 #include "FWCore/Framework/interface/stream/EDProducer.h"
0027 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0028 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0029 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0030 #include "FWCore/Utilities/interface/InputTag.h"
0031 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0032 #include "TrackingTools/PatternTools/interface/Trajectory.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 = default;
0061 
0062       static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0063         edm::ParameterSetDescription desc;
0064         desc.add<edm::InputTag>("src", edm::InputTag("generalTracks"));
0065         desc.addUntracked<bool>("copyExtras", false);
0066         desc.addUntracked<bool>("copyTrajectories", false);
0067         Selector::fillPSetDescription(desc);
0068         descriptions.addWithDefaultLabel(desc);
0069       }
0070 
0071     private:
0072       /// process one event
0073       void produce(edm::Event& evt, const edm::EventSetup& es) override {
0074         edm::Handle<reco::TrackCollection> hSrcTrack;
0075         evt.getByToken(hSrcTrackToken_, hSrcTrack);
0076 
0077         selTracks_ = std::make_unique<reco::TrackCollection>();
0078         if (copyExtras_) {
0079           selTrackExtras_ = std::make_unique<reco::TrackExtraCollection>();
0080           selHits_ = std::make_unique<TrackingRecHitCollection>();
0081         }
0082 
0083         TrackRefProd rTracks = evt.template getRefBeforePut<TrackCollection>();
0084 
0085         TrackingRecHitRefProd rHits;
0086         TrackExtraRefProd rTrackExtras;
0087         if (copyExtras_) {
0088           rHits = evt.template getRefBeforePut<TrackingRecHitCollection>();
0089           rTrackExtras = evt.template getRefBeforePut<TrackExtraCollection>();
0090         }
0091 
0092         typedef reco::TrackRef::key_type TrackRefKey;
0093         std::map<TrackRefKey, reco::TrackRef> goodTracks;
0094         TrackRefKey current = 0;
0095 
0096         selector_.init(evt, es);
0097         auto tkBegin = hSrcTrack->begin();
0098         for (reco::TrackCollection::const_iterator it = tkBegin, ed = hSrcTrack->end(); it != ed; ++it, ++current) {
0099           const reco::Track& trk = *it;
0100           const reco::TrackRef tkref(hSrcTrack, std::distance(tkBegin, it));
0101           if (!selector_(tkref))
0102             continue;
0103 
0104           selTracks_->push_back(Track(trk));  // clone and store
0105           if (!copyExtras_)
0106             continue;
0107 
0108           // TrackExtras
0109           selTrackExtras_->push_back(TrackExtra(trk.outerPosition(),
0110                                                 trk.outerMomentum(),
0111                                                 trk.outerOk(),
0112                                                 trk.innerPosition(),
0113                                                 trk.innerMomentum(),
0114                                                 trk.innerOk(),
0115                                                 trk.outerStateCovariance(),
0116                                                 trk.outerDetId(),
0117                                                 trk.innerStateCovariance(),
0118                                                 trk.innerDetId(),
0119                                                 trk.seedDirection()));
0120           selTracks_->back().setExtra(TrackExtraRef(rTrackExtras, selTrackExtras_->size() - 1));
0121           TrackExtra& tx = selTrackExtras_->back();
0122           // TrackingRecHits
0123           auto const firstHitIndex = selHits_->size();
0124           for (auto const& hit : trk.recHits())
0125             selHits_->push_back(hit->clone());
0126           tx.setHits(rHits, firstHitIndex, selHits_->size() - firstHitIndex);
0127           tx.setTrajParams(trk.extra()->trajParams(), trk.extra()->chi2sX5());
0128           assert(tx.trajParams().size() == tx.recHitsSize());
0129           if (copyTrajectories_) {
0130             goodTracks[current] = reco::TrackRef(rTracks, selTracks_->size() - 1);
0131           }
0132         }
0133         if (copyTrajectories_) {
0134           edm::Handle<std::vector<Trajectory> > hTraj;
0135           edm::Handle<TrajTrackAssociationCollection> hTTAss;
0136           evt.getByToken(hTTAssToken_, hTTAss);
0137           evt.getByToken(hTrajToken_, hTraj);
0138           edm::RefProd<std::vector<Trajectory> > TrajRefProd = evt.template getRefBeforePut<std::vector<Trajectory> >();
0139           selTrajs_ = std::make_unique<std::vector<Trajectory> >();
0140           selTTAss_ = std::make_unique<TrajTrackAssociationCollection>();
0141           for (size_t i = 0, n = hTraj->size(); i < n; ++i) {
0142             edm::Ref<std::vector<Trajectory> > trajRef(hTraj, i);
0143             TrajTrackAssociationCollection::const_iterator match = hTTAss->find(trajRef);
0144             if (match != hTTAss->end()) {
0145               const edm::Ref<reco::TrackCollection>& trkRef = match->val;
0146               TrackRefKey oldKey = trkRef.key();
0147               std::map<TrackRefKey, reco::TrackRef>::iterator getref = goodTracks.find(oldKey);
0148               if (getref != goodTracks.end()) {
0149                 // do the clone
0150                 selTrajs_->push_back(Trajectory(*trajRef));
0151                 selTTAss_->insert(edm::Ref<std::vector<Trajectory> >(TrajRefProd, selTrajs_->size() - 1),
0152                                   getref->second);
0153               }
0154             }
0155           }
0156         }
0157 
0158         evt.put(std::move(selTracks_));
0159         if (copyExtras_) {
0160           evt.put(std::move(selTrackExtras_));
0161           evt.put(std::move(selHits_));
0162           if (copyTrajectories_) {
0163             evt.put(std::move(selTrajs_));
0164             evt.put(std::move(selTTAss_));
0165           }
0166         }
0167       }
0168       /// source collection label
0169       edm::EDGetTokenT<reco::TrackCollection> hSrcTrackToken_;
0170       edm::EDGetTokenT<std::vector<Trajectory> > hTrajToken_;
0171       edm::EDGetTokenT<TrajTrackAssociationCollection> hTTAssToken_;
0172       /// copy only the tracks, not extras and rechits (for AOD)
0173       bool copyExtras_;
0174       /// copy also trajectories and trajectory->track associations
0175       bool copyTrajectories_;
0176       /// filter event
0177       Selector selector_;
0178       // some space
0179       std::unique_ptr<reco::TrackCollection> selTracks_;
0180       std::unique_ptr<reco::TrackExtraCollection> selTrackExtras_;
0181       std::unique_ptr<TrackingRecHitCollection> selHits_;
0182       std::unique_ptr<std::vector<Trajectory> > selTrajs_;
0183       std::unique_ptr<TrajTrackAssociationCollection> selTTAss_;
0184     };
0185   }  // namespace modules
0186 }  // namespace reco
0187 #endif