File indexing completed on 2025-01-08 03:35:56
0001 #ifndef RecoAlgos_TrackFullCloneSelectorBase_h
0002 #define RecoAlgos_TrackFullCloneSelectorBase_h
0003
0004
0005
0006
0007
0008
0009
0010
0011
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
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
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
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));
0105 if (!copyExtras_)
0106 continue;
0107
0108
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
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
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
0169 edm::EDGetTokenT<reco::TrackCollection> hSrcTrackToken_;
0170 edm::EDGetTokenT<std::vector<Trajectory> > hTrajToken_;
0171 edm::EDGetTokenT<TrajTrackAssociationCollection> hTTAssToken_;
0172
0173 bool copyExtras_;
0174
0175 bool copyTrajectories_;
0176
0177 Selector selector_;
0178
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 }
0186 }
0187 #endif