File indexing completed on 2024-04-06 12:28:06
0001
0002
0003
0004
0005
0006
0007
0008
0009
0010
0011 #include <utility>
0012 #include <vector>
0013 #include <memory>
0014 #include <algorithm>
0015 #include <map>
0016 #include "FWCore/Framework/interface/stream/EDProducer.h"
0017 #include "FWCore/Framework/interface/EventPrincipal.h"
0018 #include "FWCore/Framework/interface/Event.h"
0019 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0020 #include "FWCore/Utilities/interface/InputTag.h"
0021
0022 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0023 #include "DataFormats/TrackReco/interface/Track.h"
0024 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0025 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0026 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0027 #include "TrackingTools/PatternTools/interface/Trajectory.h"
0028 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0029
0030 using namespace reco;
0031
0032 class dso_hidden CosmicTrackSelector final : public edm::stream::EDProducer<> {
0033 private:
0034 public:
0035
0036 explicit CosmicTrackSelector(const edm::ParameterSet &cfg);
0037
0038 ~CosmicTrackSelector() override;
0039
0040 private:
0041 typedef math::XYZPoint Point;
0042
0043 void produce(edm::Event &evt, const edm::EventSetup &es) override;
0044
0045 bool select(const reco::BeamSpot &vertexBeamSpot, const reco::Track &tk);
0046
0047 edm::EDGetTokenT<reco::TrackCollection> src_;
0048 edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0049
0050 bool copyExtras_;
0051
0052 bool copyTrajectories_;
0053 edm::EDGetTokenT<std::vector<Trajectory>> srcTraj_;
0054 edm::EDGetTokenT<TrajTrackAssociationCollection> srcTass_;
0055
0056
0057 bool keepAllTracks_;
0058
0059 bool setQualityBit_;
0060 TrackBase::TrackQuality qualityToSet_;
0061
0062
0063 std::vector<double> res_par_;
0064 double chi2n_par_;
0065
0066
0067 double max_d0_;
0068 double max_z0_;
0069
0070 double min_pt_;
0071 double max_eta_;
0072
0073 uint32_t min_nHit_;
0074
0075 uint32_t min_nPixelHit_;
0076
0077 uint32_t min_layers_;
0078 uint32_t min_3Dlayers_;
0079 uint32_t max_lostLayers_;
0080
0081
0082 std::unique_ptr<reco::TrackCollection> selTracks_;
0083 std::unique_ptr<reco::TrackExtraCollection> selTrackExtras_;
0084 std::unique_ptr<TrackingRecHitCollection> selHits_;
0085 std::unique_ptr<std::vector<Trajectory>> selTrajs_;
0086 std::unique_ptr<std::vector<const Trajectory *>> selTrajPtrs_;
0087 std::unique_ptr<TrajTrackAssociationCollection> selTTAss_;
0088 reco::TrackRefProd rTracks_;
0089 reco::TrackExtraRefProd rTrackExtras_;
0090 TrackingRecHitRefProd rHits_;
0091 edm::RefProd<std::vector<Trajectory>> rTrajectories_;
0092 std::vector<reco::TrackRef> trackRefs_;
0093 };
0094
0095 #include <Math/DistFunc.h>
0096 #include "TMath.h"
0097
0098 CosmicTrackSelector::CosmicTrackSelector(const edm::ParameterSet &cfg)
0099 : src_(consumes<reco::TrackCollection>(cfg.getParameter<edm::InputTag>("src"))),
0100 beamspot_(consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamspot"))),
0101 copyExtras_(cfg.getUntrackedParameter<bool>("copyExtras", false)),
0102 copyTrajectories_(cfg.getUntrackedParameter<bool>("copyTrajectories", false)),
0103 keepAllTracks_(cfg.exists("keepAllTracks") ? cfg.getParameter<bool>("keepAllTracks")
0104 : false),
0105 setQualityBit_(false),
0106 qualityToSet_(TrackBase::undefQuality),
0107 chi2n_par_(cfg.getParameter<double>("chi2n_par")),
0108
0109 max_d0_(cfg.getParameter<double>("max_d0")),
0110 max_z0_(cfg.getParameter<double>("max_z0")),
0111
0112 min_pt_(cfg.getParameter<double>("min_pt")),
0113 max_eta_(cfg.getParameter<double>("max_eta")),
0114
0115 min_nHit_(cfg.getParameter<uint32_t>("min_nHit")),
0116
0117 min_nPixelHit_(cfg.getParameter<uint32_t>("min_nPixelHit")),
0118
0119 min_layers_(cfg.getParameter<uint32_t>("minNumberLayers")),
0120 min_3Dlayers_(cfg.getParameter<uint32_t>("minNumber3DLayers")),
0121 max_lostLayers_(cfg.getParameter<uint32_t>("maxNumberLostLayers")) {
0122 if (cfg.exists("qualityBit")) {
0123 std::string qualityStr = cfg.getParameter<std::string>("qualityBit");
0124 if (!qualityStr.empty()) {
0125 setQualityBit_ = true;
0126 qualityToSet_ = TrackBase::qualityByName(cfg.getParameter<std::string>("qualityBit"));
0127 }
0128 }
0129 if (keepAllTracks_ && !setQualityBit_)
0130 throw cms::Exception("Configuration")
0131 << "If you set 'keepAllTracks' to true, you must specify which qualityBit to set.\n";
0132 if (setQualityBit_ && (qualityToSet_ == TrackBase::undefQuality))
0133 throw cms::Exception("Configuration")
0134 << "You can't set the quality bit " << cfg.getParameter<std::string>("qualityBit")
0135 << " as it is 'undefQuality' or unknown.\n";
0136
0137 std::string alias(cfg.getParameter<std::string>("@module_label"));
0138 produces<reco::TrackCollection>().setBranchAlias(alias + "Tracks");
0139 if (copyExtras_) {
0140 produces<reco::TrackExtraCollection>().setBranchAlias(alias + "TrackExtras");
0141 produces<TrackingRecHitCollection>().setBranchAlias(alias + "RecHits");
0142 }
0143 if (copyTrajectories_) {
0144 srcTraj_ = consumes<std::vector<Trajectory>>(cfg.getParameter<edm::InputTag>("src"));
0145 srcTass_ = consumes<TrajTrackAssociationCollection>(cfg.getParameter<edm::InputTag>("src"));
0146 produces<std::vector<Trajectory>>().setBranchAlias(alias + "Trajectories");
0147 produces<TrajTrackAssociationCollection>().setBranchAlias(alias + "TrajectoryTrackAssociations");
0148 }
0149 }
0150
0151 CosmicTrackSelector::~CosmicTrackSelector() {}
0152
0153 void CosmicTrackSelector::produce(edm::Event &evt, const edm::EventSetup &es) {
0154 using namespace std;
0155 using namespace edm;
0156 using namespace reco;
0157
0158 Handle<TrackCollection> hSrcTrack;
0159 Handle<vector<Trajectory>> hTraj;
0160 Handle<vector<Trajectory>> hTrajP;
0161 Handle<TrajTrackAssociationCollection> hTTAss;
0162
0163
0164 edm::Handle<reco::BeamSpot> hBsp;
0165 evt.getByToken(beamspot_, hBsp);
0166 reco::BeamSpot vertexBeamSpot;
0167 vertexBeamSpot = *hBsp;
0168
0169
0170 evt.getByToken(src_, hSrcTrack);
0171
0172 selTracks_ = std::make_unique<TrackCollection>();
0173 rTracks_ = evt.getRefBeforePut<TrackCollection>();
0174 if (copyExtras_) {
0175 selTrackExtras_ = std::make_unique<TrackExtraCollection>();
0176 selHits_ = std::make_unique<TrackingRecHitCollection>();
0177 rHits_ = evt.getRefBeforePut<TrackingRecHitCollection>();
0178 rTrackExtras_ = evt.getRefBeforePut<TrackExtraCollection>();
0179 }
0180
0181 if (copyTrajectories_)
0182 trackRefs_.resize(hSrcTrack->size());
0183
0184
0185 size_t current = 0;
0186 for (TrackCollection::const_iterator it = hSrcTrack->begin(), ed = hSrcTrack->end(); it != ed; ++it, ++current) {
0187 const Track &trk = *it;
0188
0189 bool ok = select(vertexBeamSpot, trk);
0190 if (!ok) {
0191 if (copyTrajectories_)
0192 trackRefs_[current] = reco::TrackRef();
0193 if (!keepAllTracks_)
0194 continue;
0195 }
0196 selTracks_->push_back(Track(trk));
0197 if (ok && setQualityBit_)
0198 selTracks_->back().setQuality(qualityToSet_);
0199 if (copyExtras_) {
0200
0201 selTrackExtras_->push_back(TrackExtra(trk.outerPosition(),
0202 trk.outerMomentum(),
0203 trk.outerOk(),
0204 trk.innerPosition(),
0205 trk.innerMomentum(),
0206 trk.innerOk(),
0207 trk.outerStateCovariance(),
0208 trk.outerDetId(),
0209 trk.innerStateCovariance(),
0210 trk.innerDetId(),
0211 trk.seedDirection(),
0212 trk.seedRef()));
0213 selTracks_->back().setExtra(TrackExtraRef(rTrackExtras_, selTrackExtras_->size() - 1));
0214 TrackExtra &tx = selTrackExtras_->back();
0215 tx.setResiduals(trk.residuals());
0216
0217 auto const firstHitIndex = selHits_->size();
0218 for (trackingRecHit_iterator hit = trk.recHitsBegin(); hit != trk.recHitsEnd(); ++hit) {
0219 selHits_->push_back((*hit)->clone());
0220 }
0221 tx.setHits(rHits_, firstHitIndex, selHits_->size() - firstHitIndex);
0222 tx.setTrajParams(trk.extra()->trajParams(), trk.extra()->chi2sX5());
0223 }
0224 if (copyTrajectories_) {
0225 trackRefs_[current] = TrackRef(rTracks_, selTracks_->size() - 1);
0226 }
0227 }
0228 if (copyTrajectories_) {
0229 Handle<vector<Trajectory>> hTraj;
0230 Handle<TrajTrackAssociationCollection> hTTAss;
0231 evt.getByToken(srcTass_, hTTAss);
0232 evt.getByToken(srcTraj_, hTraj);
0233 selTrajs_ = std::make_unique<std::vector<Trajectory>>();
0234 rTrajectories_ = evt.getRefBeforePut<vector<Trajectory>>();
0235 selTTAss_ = std::make_unique<TrajTrackAssociationCollection>(&evt.productGetter());
0236 for (size_t i = 0, n = hTraj->size(); i < n; ++i) {
0237 Ref<vector<Trajectory>> trajRef(hTraj, i);
0238 TrajTrackAssociationCollection::const_iterator match = hTTAss->find(trajRef);
0239 if (match != hTTAss->end()) {
0240 const Ref<TrackCollection> &trkRef = match->val;
0241 short oldKey = static_cast<short>(trkRef.key());
0242 if (trackRefs_[oldKey].isNonnull()) {
0243 selTrajs_->push_back(Trajectory(*trajRef));
0244 selTTAss_->insert(Ref<vector<Trajectory>>(rTrajectories_, selTrajs_->size() - 1), trackRefs_[oldKey]);
0245 }
0246 }
0247 }
0248 }
0249
0250 static const std::string emptyString;
0251 evt.put(std::move(selTracks_));
0252 if (copyExtras_) {
0253 evt.put(std::move(selTrackExtras_));
0254 evt.put(std::move(selHits_));
0255 }
0256 if (copyTrajectories_) {
0257 evt.put(std::move(selTrajs_));
0258 evt.put(std::move(selTTAss_));
0259 }
0260 }
0261
0262 bool CosmicTrackSelector::select(const reco::BeamSpot &vertexBeamSpot, const reco::Track &tk) {
0263
0264
0265 using namespace std;
0266
0267
0268 uint32_t nlayers = tk.hitPattern().trackerLayersWithMeasurement();
0269 uint32_t nlayers3D =
0270 tk.hitPattern().pixelLayersWithMeasurement() + tk.hitPattern().numberOfValidStripLayersWithMonoAndStereo();
0271 uint32_t nlayersLost = tk.hitPattern().trackerLayersWithoutMeasurement(reco::HitPattern::TRACK_HITS);
0272
0273
0274 uint32_t nHit = 0;
0275 uint32_t nPixelHit = 0;
0276 for (trackingRecHit_iterator recHit = tk.recHitsBegin(); recHit != tk.recHitsEnd(); ++recHit) {
0277 if (!((*recHit)->isValid()))
0278 continue;
0279 ++nHit;
0280 DetId id((*recHit)->geographicalId());
0281 if ((unsigned int)id.subdetId() == PixelSubdetector::PixelBarrel ||
0282 (unsigned int)id.subdetId() == PixelSubdetector::PixelEndcap)
0283 ++nPixelHit;
0284 }
0285
0286
0287 if (nHit < min_nHit_)
0288 return false;
0289
0290 if (nPixelHit < min_nPixelHit_)
0291 return false;
0292 if (nlayers < min_layers_)
0293 return false;
0294 if (nlayers3D < min_3Dlayers_)
0295 return false;
0296 if (nlayersLost > max_lostLayers_)
0297 return false;
0298
0299
0300 double pt = tk.pt(), eta = tk.eta(), chi2n = tk.normalizedChi2();
0301 double d0 = -tk.dxy(vertexBeamSpot.position()), dz = tk.dz();
0302
0303
0304 if (abs(d0) > max_d0_)
0305 return false;
0306 if (abs(dz) > max_z0_)
0307 return false;
0308
0309
0310 if (abs(eta) > max_eta_)
0311 return false;
0312 if (pt < min_pt_)
0313 return false;
0314 if (chi2n > chi2n_par_ * nlayers)
0315 return false;
0316
0317 else
0318 return true;
0319 }
0320
0321 #include "FWCore/PluginManager/interface/ModuleDef.h"
0322 #include "FWCore/Framework/interface/MakerMacros.h"
0323
0324 DEFINE_FWK_MODULE(CosmicTrackSelector);