FinalTrackSelectors/plugins/CosmicTrackSelector.cc

0001 /** \class CosmicTrackSelector
0002  *
0003  * selects a subset of a track collection, copying extra information on demand
0004  *
0005  * \author Paolo Azzurri, Giovanni Petrucciani
0006  *
0007  *
0008  *
0009  */
0010
0011 #include <algorithm>
0012 #include <map>
0013 #include <memory>
0014 #include <utility>
0015 #include <vector>
0016
0017 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0018 #include "DataFormats/TrackReco/interface/Track.h"
0019 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0020 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0021 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0022 #include "FWCore/Framework/interface/Event.h"
0023 #include "FWCore/Framework/interface/EventPrincipal.h"
0024 #include "FWCore/Framework/interface/stream/EDProducer.h"
0025 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0026 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0027 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0028 #include "FWCore/Utilities/interface/InputTag.h"
0029 #include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"
0030 #include "TrackingTools/PatternTools/interface/Trajectory.h"
0031
0032 #include <Math/DistFunc.h>
0033 #include "TMath.h"
0034
0035 using namespace reco;
0036
0037 class dso_hidden CosmicTrackSelector final : public edm::stream::EDProducer<> {
0038 private:
0039 public:
0040   // constructor
0041   explicit CosmicTrackSelector(const edm::ParameterSet &cfg);
0042   // destructor
0043   ~CosmicTrackSelector() override = default;
0044
0045   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0046
0047 private:
0048   typedef math::XYZPoint Point;
0049   // process one event
0050   void produce(edm::Event &evt, const edm::EventSetup &es) override;
0051   // return class, or -1 if rejected
0052   bool select(const reco::BeamSpot &vertexBeamSpot, const reco::Track &tk);
0053   // source collection label
0054   edm::EDGetTokenT<reco::TrackCollection> src_;
0055   edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0056   // copy only the tracks, not extras and rechits (for AOD)
0057   bool copyExtras_;
0058   // copy also trajectories and trajectory->track associations
0059   bool copyTrajectories_;
0060   edm::EDGetTokenT<std::vector<Trajectory>> srcTraj_;
0061   edm::EDGetTokenT<TrajTrackAssociationCollection> srcTass_;
0062
0063   // save all the tracks
0064   bool keepAllTracks_;
0065   // do I have to set a quality bit?
0066   bool setQualityBit_;
0067   TrackBase::TrackQuality qualityToSet_;
0068
0069   //  parameters for adapted optimal cuts on chi2 and primary vertex compatibility
0070   std::vector<double> res_par_;
0071   double chi2n_par_;
0072
0073   // Impact parameter absolute cuts
0074   double max_d0_;
0075   double max_z0_;
0076   // Trackk parameter cuts
0077   double min_pt_;
0078   double max_eta_;
0079   // Cut on number of valid hits
0080   uint32_t min_nHit_;
0081   // Cut on number of valid Pixel hits
0082   uint32_t min_nPixelHit_;
0083   // Cuts on numbers of layers with hits/3D hits/lost hits.
0084   uint32_t min_layers_;
0085   uint32_t min_3Dlayers_;
0086   uint32_t max_lostLayers_;
0087
0088   // storage
0089   std::unique_ptr<reco::TrackCollection> selTracks_;
0090   std::unique_ptr<reco::TrackExtraCollection> selTrackExtras_;
0091   std::unique_ptr<TrackingRecHitCollection> selHits_;
0092   std::unique_ptr<std::vector<Trajectory>> selTrajs_;
0093   std::unique_ptr<std::vector<const Trajectory *>> selTrajPtrs_;
0094   std::unique_ptr<TrajTrackAssociationCollection> selTTAss_;
0095   reco::TrackRefProd rTracks_;
0096   reco::TrackExtraRefProd rTrackExtras_;
0097   TrackingRecHitRefProd rHits_;
0098   edm::RefProd<std::vector<Trajectory>> rTrajectories_;
0099   std::vector<reco::TrackRef> trackRefs_;
0100 };
0101
0102 CosmicTrackSelector::CosmicTrackSelector(const edm::ParameterSet &cfg)
0103     : src_(consumes<reco::TrackCollection>(cfg.getParameter<edm::InputTag>("src"))),
0104       beamspot_(consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamspot"))),
0105       copyExtras_(cfg.getUntrackedParameter<bool>("copyExtras", false)),
0106       copyTrajectories_(cfg.getUntrackedParameter<bool>("copyTrajectories", false)),
0107       keepAllTracks_(cfg.getParameter<bool>("keepAllTracks")),
0108       setQualityBit_(false),
0109       qualityToSet_(TrackBase::undefQuality),
0110       chi2n_par_(cfg.getParameter<double>("chi2n_par")),
0111       // Impact parameter absolute cuts.
0112       max_d0_(cfg.getParameter<double>("max_d0")),
0113       max_z0_(cfg.getParameter<double>("max_z0")),
0114       // Track parameter cuts.
0115       min_pt_(cfg.getParameter<double>("min_pt")),
0116       max_eta_(cfg.getParameter<double>("max_eta")),
0117       // Cut on number of valid hits
0118       min_nHit_(cfg.getParameter<uint32_t>("min_nHit")),
0119       // Cut on number of valid hits
0120       min_nPixelHit_(cfg.getParameter<uint32_t>("min_nPixelHit")),
0121       // Cuts on numbers of layers with hits/3D hits/lost hits.
0122       min_layers_(cfg.getParameter<uint32_t>("minNumberLayers")),
0123       min_3Dlayers_(cfg.getParameter<uint32_t>("minNumber3DLayers")),
0124       max_lostLayers_(cfg.getParameter<uint32_t>("maxNumberLostLayers")) {
0125   std::string qualityStr = cfg.getParameter<std::string>("qualityBit");
0126   if (!qualityStr.empty()) {
0127     setQualityBit_ = true;
0128     qualityToSet_ = TrackBase::qualityByName(cfg.getParameter<std::string>("qualityBit"));
0129   }
0130
0131   if (keepAllTracks_ && !setQualityBit_)
0132     throw cms::Exception("Configuration")
0133         << "If you set 'keepAllTracks' to true, you must specify which qualityBit to set.\n";
0134   if (setQualityBit_ && (qualityToSet_ == TrackBase::undefQuality))
0135     throw cms::Exception("Configuration")
0136         << "You can't set the quality bit " << cfg.getParameter<std::string>("qualityBit")
0137         << " as it is 'undefQuality' or unknown.\n";
0138
0139   std::string alias(cfg.getParameter<std::string>("@module_label"));
0140   produces<reco::TrackCollection>().setBranchAlias(alias + "Tracks");
0141   if (copyExtras_) {
0142     produces<reco::TrackExtraCollection>().setBranchAlias(alias + "TrackExtras");
0143     produces<TrackingRecHitCollection>().setBranchAlias(alias + "RecHits");
0144   }
0145   if (copyTrajectories_) {
0146     srcTraj_ = consumes<std::vector<Trajectory>>(cfg.getParameter<edm::InputTag>("src"));
0147     srcTass_ = consumes<TrajTrackAssociationCollection>(cfg.getParameter<edm::InputTag>("src"));
0148     produces<std::vector<Trajectory>>().setBranchAlias(alias + "Trajectories");
0149     produces<TrajTrackAssociationCollection>().setBranchAlias(alias + "TrajectoryTrackAssociations");
0150   }
0151 }
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   // looking for the beam spot
0164   edm::Handle<reco::BeamSpot> hBsp;
0165   evt.getByToken(beamspot_, hBsp);
0166   reco::BeamSpot vertexBeamSpot;
0167   vertexBeamSpot = *hBsp;
0168
0169   // Get tracks
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   // Loop over tracks
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     // Check if this track passes cuts
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));  // clone and store
0197     if (ok && setQualityBit_)
0198       selTracks_->back().setQuality(qualityToSet_);
0199     if (copyExtras_) {
0200       // TrackExtras
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       // TrackingRecHits
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   // Decide if the given track passes selection cuts.
0264
0265   using namespace std;
0266
0267   // Cuts on numbers of layers with hits/3D hits/lost hits.
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   // Get the number of valid hits and PixelHits
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   // Cut on the number of valid hits
0287   if (nHit < min_nHit_)
0288     return false;
0289   // Cut on the number of valid Pixel hits
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   // Get track parameters
0300   double pt = tk.pt(), eta = tk.eta(), chi2n = tk.normalizedChi2();
0301   double d0 = -tk.dxy(vertexBeamSpot.position()), dz = tk.dz();
0302
0303   // Absolute cuts on all tracks impact parameters with respect to beam-spot.
0304   if (abs(d0) > max_d0_)
0305     return false;
0306   if (abs(dz) > max_z0_)
0307     return false;
0308
0309   // optimized cuts adapted to the track eta, pt and  chiquare/ndof
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 void CosmicTrackSelector::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0322   edm::ParameterSetDescription desc;
0323   desc.add<edm::InputTag>("src", edm::InputTag("ctfWithMaterialTracksCosmics"));
0324   desc.add<edm::InputTag>("beamspot", edm::InputTag("offlineBeamSpot"));
0325   desc.addUntracked<bool>("copyExtras", true);
0326   desc.addUntracked<bool>("copyTrajectories", false);
0327   desc.add<bool>("keepAllTracks", false)
0328       ->setComment("if set to true tracks failing this filter are kept in the output");
0329   desc.add<double>("chi2n_par", 10.0)->setComment("parameters for adapted optimal cuts on chi2");
0330   desc.add<double>("max_d0", 110.)->setComment("d0 impact parameter absolute cut");
0331   desc.add<double>("max_z0", 300.)->setComment("z0 impact parameter absolute cut");
0332   desc.add<double>("min_pt", 1.0);
0333   desc.add<double>("max_eta", 2.0);
0334   desc.add<uint32_t>("min_nHit", 5);
0335   desc.add<uint32_t>("min_nPixelHit", 0);
0336   desc.add<uint32_t>("minNumberLayers", 0);
0337   desc.add<uint32_t>("minNumber3DLayers", 0);
0338   desc.add<uint32_t>("maxNumberLostLayers", 999);
0339   desc.add<std::string>("qualityBit", std::string(""))
0340       ->setComment("set to '' or comment out if you don't want to set the bit");
0341   descriptions.addWithDefaultLabel(desc);
0342 }
0343
0344 #include "FWCore/PluginManager/interface/ModuleDef.h"
0345 #include "FWCore/Framework/interface/MakerMacros.h"
0346
0347 DEFINE_FWK_MODULE(CosmicTrackSelector);