00001 #include "RecoTracker/FinalTrackSelectors/interface/TrackMultiSelector.h"
00002 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
00003
00004 #include <Math/DistFunc.h>
00005 #include "TMath.h"
00006
00007 using reco::modules::TrackMultiSelector;
00008
00009
00010 TrackMultiSelector::Block::Block(const edm::ParameterSet & cfg) :
00011 pt(p2p<double>(cfg,"pt")),
00012 vlayers(p2p<uint32_t>(cfg,"validLayers")),
00013 lhits(p2p<uint32_t>(cfg,"lostHits")),
00014 chi2n(p2p<double>(cfg,"chi2n")),
00015 d0(cfg.getParameter<double>("d0")),
00016 dz(cfg.getParameter<double>("dz")),
00017 d0Rel(cfg.getParameter<double>("d0Rel")),
00018 dzRel(cfg.getParameter<double>("dzRel"))
00019 {
00020 }
00021
00022 TrackMultiSelector::TrackMultiSelector( const edm::ParameterSet & cfg ) :
00023 src_( cfg.getParameter<edm::InputTag>( "src" ) ),
00024 vertices_( cfg.getParameter<edm::InputTag>( "vertices" ) ),
00025 copyExtras_(cfg.getUntrackedParameter<bool>("copyExtras", false)),
00026 copyTrajectories_(cfg.getUntrackedParameter<bool>("copyTrajectories", false)),
00027 splitOutputs_( cfg.getUntrackedParameter<bool>("splitOutputs", false) ),
00028 vtxNumber_( cfg.getParameter<int32_t>("vtxNumber") ),
00029 vtxTracks_( cfg.getParameter<uint32_t>("vtxTracks") ),
00030 vtxChi2Prob_( cfg.getParameter<double>("vtxChi2Prob") )
00031 {
00032 edm::ParameterSet beamSpotPSet = cfg.getParameter<edm::ParameterSet>("beamspot");
00033 beamspot_ = beamSpotPSet.getParameter<edm::InputTag>("src");
00034 beamspotDZsigmas_ = beamSpotPSet.getParameter<double>("dzSigmas");
00035 beamspotD0_ = beamSpotPSet.getParameter<double>("d0");
00036
00037 typedef std::vector<edm::ParameterSet> VPSet;
00038 VPSet psets = cfg.getParameter<VPSet>("cutSets");
00039 blocks_.reserve(psets.size());
00040 for (VPSet::const_iterator it = psets.begin(), ed = psets.end(); it != ed; ++it) {
00041 blocks_.push_back(TrackMultiSelector::Block(*it));
00042 }
00043
00044 if (splitOutputs_) {
00045 char buff[15];
00046 for (size_t i = 0; i < blocks_.size(); ++i) {
00047 sprintf(buff,"set%d", static_cast<int>(i+1));
00048 labels_.push_back(std::string(buff));
00049 }
00050 } else {
00051 labels_.push_back(std::string(""));
00052 }
00053
00054 std::string alias( cfg.getParameter<std::string>( "@module_label" ) );
00055 for (std::vector<std::string>::const_iterator li = labels_.begin(), le = labels_.end(); li != le; ++li) {
00056 const char *l= li->c_str();
00057 produces<reco::TrackCollection>(l).setBranchAlias( alias + "Tracks" + l);
00058 if (copyExtras_) {
00059 produces<reco::TrackExtraCollection>(l).setBranchAlias( alias + "TrackExtras" + l);
00060 produces<TrackingRecHitCollection>(l).setBranchAlias( alias + "RecHits" + l);
00061 if (copyTrajectories_) {
00062 produces< std::vector<Trajectory> >(l).setBranchAlias( alias + "Trajectories" + l);
00063 produces< TrajTrackAssociationCollection >(l).setBranchAlias( alias + "TrajectoryTrackAssociations" + l);
00064 }
00065 }
00066 }
00067
00068 size_t nblocks = splitOutputs_ ? blocks_.size() : 1;
00069 selTracks_ = new std::auto_ptr<reco::TrackCollection>[nblocks];
00070 selTrackExtras_ = new std::auto_ptr<reco::TrackExtraCollection>[nblocks];
00071 selHits_ = new std::auto_ptr<TrackingRecHitCollection>[nblocks];
00072 selTrajs_ = new std::auto_ptr< std::vector<Trajectory> >[nblocks];
00073 selTTAss_ = new std::auto_ptr< TrajTrackAssociationCollection >[nblocks];
00074 rTracks_ = std::vector<reco::TrackRefProd>(nblocks);
00075 rHits_ = std::vector<TrackingRecHitRefProd>(nblocks);
00076 rTrackExtras_ = std::vector<reco::TrackExtraRefProd>(nblocks);
00077 rTrajectories_ = std::vector< edm::RefProd< std::vector<Trajectory> > >(nblocks);
00078 for (size_t i = 0; i < nblocks; i++) {
00079 selTracks_[i] = std::auto_ptr<reco::TrackCollection>(new reco::TrackCollection());
00080 selTrackExtras_[i] = std::auto_ptr<reco::TrackExtraCollection>(new reco::TrackExtraCollection());
00081 selHits_[i] = std::auto_ptr<TrackingRecHitCollection>(new TrackingRecHitCollection());
00082 selTrajs_[i] = std::auto_ptr< std::vector<Trajectory> >(new std::vector<Trajectory>());
00083 selTTAss_[i] = std::auto_ptr< TrajTrackAssociationCollection >(new TrajTrackAssociationCollection());
00084 }
00085
00086 }
00087
00088 TrackMultiSelector::~TrackMultiSelector() {
00089 delete [] selTracks_;
00090 delete [] selTrackExtras_;
00091 delete [] selHits_;
00092 delete [] selTrajs_;
00093 delete [] selTTAss_;
00094 }
00095
00096 void TrackMultiSelector::produce( edm::Event& evt, const edm::EventSetup& es )
00097 {
00098 using namespace std;
00099 using namespace edm;
00100 using namespace reco;
00101
00102 size_t nblocks = splitOutputs_ ? blocks_.size() : 1;
00103
00104 Handle<TrackCollection> hSrcTrack;
00105 Handle< vector<Trajectory> > hTraj;
00106 Handle< TrajTrackAssociationCollection > hTTAss;
00107
00108 edm::Handle<reco::VertexCollection> hVtx;
00109 evt.getByLabel(vertices_, hVtx);
00110 std::vector<Point> points;
00111 if (vtxNumber_ != 0) selectVertices(*hVtx, points);
00112
00113 edm::Handle<reco::BeamSpot> hBsp;
00114 evt.getByLabel(beamspot_, hBsp);
00115
00116 evt.getByLabel( src_, hSrcTrack );
00117
00118 for (size_t i = 0; i < nblocks; i++) {
00119 selTracks_[i] = auto_ptr<TrackCollection>(new TrackCollection());
00120 rTracks_[i] = evt.getRefBeforePut<TrackCollection>(labels_[i]);
00121 if (copyExtras_) {
00122 selTrackExtras_[i] = auto_ptr<TrackExtraCollection>(new TrackExtraCollection());
00123 selHits_[i] = auto_ptr<TrackingRecHitCollection>(new TrackingRecHitCollection());
00124 rHits_[i] = evt.getRefBeforePut<TrackingRecHitCollection>(labels_[i]);
00125 rTrackExtras_[i] = evt.getRefBeforePut<TrackExtraCollection>(labels_[i]);
00126 }
00127 }
00128
00129 if (copyTrajectories_) whereItWent_.resize(hSrcTrack->size());
00130 size_t current = 0;
00131 for (TrackCollection::const_iterator it = hSrcTrack->begin(), ed = hSrcTrack->end(); it != ed; ++it, ++current) {
00132 const Track & trk = * it;
00133 short where = select(trk, *hBsp, points);
00134 if (where == -1) {
00135 if (copyTrajectories_) whereItWent_[current] = std::pair<short, reco::TrackRef>(-1, reco::TrackRef());
00136 continue;
00137 }
00138 if (!splitOutputs_) where = 0;
00139 selTracks_[where]->push_back( Track( trk ) );
00140 if (!copyExtras_) continue;
00141
00142
00143 selTrackExtras_[where]->push_back( TrackExtra( trk.outerPosition(), trk.outerMomentum(), trk.outerOk(),
00144 trk.innerPosition(), trk.innerMomentum(), trk.innerOk(),
00145 trk.outerStateCovariance(), trk.outerDetId(),
00146 trk.innerStateCovariance(), trk.innerDetId(),
00147 trk.seedDirection(), trk.seedRef() ) );
00148 selTracks_[where]->back().setExtra( TrackExtraRef( rTrackExtras_[where], selTrackExtras_[where]->size() - 1) );
00149 TrackExtra & tx = selTrackExtras_[where]->back();
00150
00151 for( trackingRecHit_iterator hit = trk.recHitsBegin(); hit != trk.recHitsEnd(); ++ hit ) {
00152 selHits_[where]->push_back( (*hit)->clone() );
00153 tx.add( TrackingRecHitRef( rHits_[where], selHits_[where]->size() - 1) );
00154 }
00155 if (copyTrajectories_) {
00156 whereItWent_[current] = std::pair<short, reco::TrackRef>(where, TrackRef(rTracks_[where], selTracks_[where]->size() - 1));
00157 }
00158 }
00159 if ( copyTrajectories_ ) {
00160 Handle< vector<Trajectory> > hTraj;
00161 Handle< TrajTrackAssociationCollection > hTTAss;
00162 evt.getByLabel(src_, hTTAss);
00163 evt.getByLabel(src_, hTraj);
00164 for (size_t i = 0; i < nblocks; i++) {
00165 rTrajectories_[i] = evt.getRefBeforePut< vector<Trajectory> >(labels_[i]);
00166 selTrajs_[i] = auto_ptr< vector<Trajectory> >(new vector<Trajectory>());
00167 selTTAss_[i] = auto_ptr< TrajTrackAssociationCollection >(new TrajTrackAssociationCollection());
00168 }
00169 for (size_t i = 0, n = hTraj->size(); i < n; ++i) {
00170 Ref< vector<Trajectory> > trajRef(hTraj, i);
00171 TrajTrackAssociationCollection::const_iterator match = hTTAss->find(trajRef);
00172 if (match != hTTAss->end()) {
00173 const Ref<TrackCollection> &trkRef = match->val;
00174 short oldKey = static_cast<short>(trkRef.key());
00175 if (whereItWent_[oldKey].first != -1) {
00176 int where = whereItWent_[oldKey].first;
00177 selTrajs_[where]->push_back( Trajectory(*trajRef) );
00178 selTTAss_[where]->insert ( Ref< vector<Trajectory> >(rTrajectories_[where], selTrajs_[where]->size() - 1), whereItWent_[oldKey].second );
00179 }
00180 }
00181 }
00182 }
00183
00184
00185 static const std::string emptyString;
00186 for (size_t i = 0; i < nblocks; i++) {
00187 const std::string & lbl = ( splitOutputs_ ? labels_[i] : emptyString);
00188 evt.put(selTracks_[i], lbl);
00189 if (copyExtras_ ) {
00190 evt.put(selTrackExtras_[i], lbl);
00191 evt.put(selHits_[i], lbl);
00192 if ( copyTrajectories_ ) {
00193 evt.put(selTrajs_[i], lbl);
00194 evt.put(selTTAss_[i], lbl);
00195 }
00196 }
00197 }
00198 }
00199
00200 inline bool TrackMultiSelector::testVtx ( const reco::Track &tk, const reco::BeamSpot &beamSpot,
00201 const std::vector<Point> &points,
00202 const TrackMultiSelector::Block &cut) {
00203 using std::abs;
00204 double d0Err =abs(tk.d0Error()), dzErr = abs(tk.dzError());
00205 if (points.empty()) {
00206 Point spot = beamSpot.position();
00207 double dz = abs(tk.dz(spot)), d0 = abs(tk.dxy(spot));
00208 return ( dz < beamspotDZsigmas_*beamSpot.sigmaZ() ) && ( d0 < beamspotD0_ );
00209 }
00210 for (std::vector<Point>::const_iterator point = points.begin(), end = points.end(); point != end; ++point) {
00211 double dz = abs(tk.dz(*point)), d0 = abs(tk.dxy(*point));
00212 if ((dz < cut.dz) && (d0 < cut.d0)
00213 && fabs(dz/std::max(dzErr,1e-9)) < cut.dzRel && (d0/std::max(d0Err,1e-8) < cut.d0Rel )) return true;
00214 }
00215 return false;
00216 }
00217
00218 short TrackMultiSelector::select(const reco::Track &tk, const reco::BeamSpot &beamSpot, const std::vector<Point> &points) {
00219 uint32_t vlayers = tk.hitPattern().trackerLayersWithMeasurement(), lhits = tk.numberOfLostHits();
00220 double pt = tk.pt(), chi2n = tk.normalizedChi2();
00221 int which = 0;
00222 for (std::vector<TrackMultiSelector::Block>::const_iterator itb = blocks_.begin(), edb = blocks_.end(); itb != edb; ++itb, ++which) {
00223 if ( ( itb->vlayers.first <= vlayers ) && ( vlayers <= itb->vlayers.second ) &&
00224 ( itb->chi2n.first <= chi2n ) && ( chi2n <= itb->chi2n.second ) &&
00225 ( itb->pt.first <= pt ) && ( pt <= itb->pt.second ) &&
00226 ( itb->lhits.first <= lhits ) && ( lhits <= itb->lhits.second ) &&
00227 testVtx(tk, beamSpot, points, *itb) )
00228 {
00229 return which;
00230 }
00231 }
00232 return -1;
00233 }
00234 void TrackMultiSelector::selectVertices(const reco::VertexCollection &vtxs, std::vector<Point> &points) {
00235 using namespace reco;
00236
00237 int32_t toTake = vtxNumber_;
00238 for (VertexCollection::const_iterator it = vtxs.begin(), ed = vtxs.end(); it != ed; ++it) {
00239 if ((it->tracksSize() >= vtxTracks_) &&
00240 ( (it->chi2() == 0.0) || (TMath::Prob(it->chi2(), static_cast<int32_t>(it->ndof()) ) >= vtxChi2Prob_) ) ) {
00241 points.push_back(it->position());
00242 toTake--; if (toTake == 0) break;
00243 }
00244 }
00245 }
1.5.4