Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2024-06-25 22:35:16 

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/Framework/interface/EventSetup.h"
0005 #include "FWCore/Framework/interface/ESHandle.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0007 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0008 
0009 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0010 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0011 #include "DataFormats/Common/interface/ValueMap.h"
0012 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0013 #include "DataFormats/TrackReco/interface/Track.h"
0014 #include "DataFormats/VertexReco/interface/Vertex.h"
0015 #include <CLHEP/Units/GlobalPhysicalConstants.h>
0016 #include <CLHEP/Units/SystemOfUnits.h>
0017 
0018 using namespace std;
0019 using namespace edm;
0020 
0021 class TOFPIDProducer : public edm::stream::EDProducer<> {
0022 public:
0023   TOFPIDProducer(const ParameterSet& pset);
0024 
0025   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0026 
0027   template <class H, class T>
0028   void fillValueMap(edm::Event& iEvent,
0029                     const edm::Handle<H>& handle,
0030                     const std::vector<T>& vec,
0031                     const std::string& name) const;
0032 
0033   void produce(edm::Event& ev, const edm::EventSetup& es) final;
0034 
0035 private:
0036   static constexpr char t0Name[] = "t0";
0037   static constexpr char sigmat0Name[] = "sigmat0";
0038   static constexpr char t0safeName[] = "t0safe";
0039   static constexpr char sigmat0safeName[] = "sigmat0safe";
0040   static constexpr char probPiName[] = "probPi";
0041   static constexpr char probKName[] = "probK";
0042   static constexpr char probPName[] = "probP";
0043 
0044   edm::EDGetTokenT<reco::TrackCollection> tracksToken_;
0045   edm::EDGetTokenT<edm::ValueMap<float>> t0Token_;
0046   edm::EDGetTokenT<edm::ValueMap<float>> tmtdToken_;
0047   edm::EDGetTokenT<edm::ValueMap<float>> sigmat0Token_;
0048   edm::EDGetTokenT<edm::ValueMap<float>> sigmatmtdToken_;
0049   edm::EDGetTokenT<edm::ValueMap<float>> tofkToken_;
0050   edm::EDGetTokenT<edm::ValueMap<float>> tofpToken_;
0051   edm::EDGetTokenT<edm::ValueMap<float>> sigmatofpiToken_;
0052   edm::EDGetTokenT<edm::ValueMap<float>> sigmatofkToken_;
0053   edm::EDGetTokenT<edm::ValueMap<float>> sigmatofpToken_;
0054   edm::EDGetTokenT<reco::VertexCollection> vtxsToken_;
0055   edm::EDGetTokenT<edm::ValueMap<float>> trackMTDTimeQualityToken_;
0056   const double vtxMaxSigmaT_;
0057   const double maxDz_;
0058   const double maxDtSignificance_;
0059   const double minProbHeavy_;
0060   const double fixedT0Error_;
0061   const double probPion_;
0062   const double probKaon_;
0063   const double probProton_;
0064   const double minTrackTimeQuality_;
0065   const bool MVASel_;
0066   const bool vertexReassignment_;
0067 };
0068 
0069 TOFPIDProducer::TOFPIDProducer(const ParameterSet& iConfig)
0070     : tracksToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("tracksSrc"))),
0071       t0Token_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("t0Src"))),
0072       tmtdToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("tmtdSrc"))),
0073       sigmat0Token_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmat0Src"))),
0074       sigmatmtdToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmatmtdSrc"))),
0075       tofkToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("tofkSrc"))),
0076       tofpToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("tofpSrc"))),
0077       sigmatofpiToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmatofpiSrc"))),
0078       sigmatofkToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmatofkSrc"))),
0079       sigmatofpToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmatofpSrc"))),
0080       vtxsToken_(consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("vtxsSrc"))),
0081       trackMTDTimeQualityToken_(
0082           consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("trackMTDTimeQualityVMapTag"))),
0083       vtxMaxSigmaT_(iConfig.getParameter<double>("vtxMaxSigmaT")),
0084       maxDz_(iConfig.getParameter<double>("maxDz")),
0085       maxDtSignificance_(iConfig.getParameter<double>("maxDtSignificance")),
0086       minProbHeavy_(iConfig.getParameter<double>("minProbHeavy")),
0087       fixedT0Error_(iConfig.getParameter<double>("fixedT0Error")),
0088       probPion_(iConfig.getParameter<double>("probPion")),
0089       probKaon_(iConfig.getParameter<double>("probKaon")),
0090       probProton_(iConfig.getParameter<double>("probProton")),
0091       minTrackTimeQuality_(iConfig.getParameter<double>("minTrackTimeQuality")),
0092       MVASel_(iConfig.getParameter<bool>("MVASel")),
0093       vertexReassignment_(iConfig.getParameter<bool>("vertexReassignment")) {
0094   produces<edm::ValueMap<float>>(t0Name);
0095   produces<edm::ValueMap<float>>(sigmat0Name);
0096   produces<edm::ValueMap<float>>(t0safeName);
0097   produces<edm::ValueMap<float>>(sigmat0safeName);
0098   produces<edm::ValueMap<float>>(probPiName);
0099   produces<edm::ValueMap<float>>(probKName);
0100   produces<edm::ValueMap<float>>(probPName);
0101 }
0102 
0103 // Configuration descriptions
0104 void TOFPIDProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0105   edm::ParameterSetDescription desc;
0106   desc.add<edm::InputTag>("tracksSrc", edm::InputTag("generalTracks"))->setComment("Input tracks collection");
0107   desc.add<edm::InputTag>("t0Src", edm::InputTag("trackExtenderWithMTD:generalTrackt0"))
0108       ->setComment("Input ValueMap for track time at beamline");
0109   desc.add<edm::InputTag>("tmtdSrc", edm::InputTag("trackExtenderWithMTD:generalTracktmtd"))
0110       ->setComment("Input ValueMap for track time at MTD");
0111   desc.add<edm::InputTag>("sigmat0Src", edm::InputTag("trackExtenderWithMTD:generalTracksigmat0"))
0112       ->setComment("Input ValueMap for track time uncertainty at beamline");
0113   desc.add<edm::InputTag>("sigmatmtdSrc", edm::InputTag("trackExtenderWithMTD:generalTracksigmatmtd"))
0114       ->setComment("Input ValueMap for track time uncertainty at MTD");
0115   desc.add<edm::InputTag>("tofkSrc", edm::InputTag("trackExtenderWithMTD:generalTrackTofK"))
0116       ->setComment("Input ValueMap for track tof as kaon");
0117   desc.add<edm::InputTag>("tofpSrc", edm::InputTag("trackExtenderWithMTD:generalTrackTofP"))
0118       ->setComment("Input ValueMap for track tof as proton");
0119   desc.add<edm::InputTag>("sigmatofpiSrc", edm::InputTag("trackExtenderWithMTD:generalTrackSigmaTofPi"))
0120       ->setComment("Input ValueMap for track sigma(tof) as pion");
0121   desc.add<edm::InputTag>("sigmatofkSrc", edm::InputTag("trackExtenderWithMTD:generalTrackSigmaTofK"))
0122       ->setComment("Input ValueMap for track sigma(tof) as kaon");
0123   desc.add<edm::InputTag>("sigmatofpSrc", edm::InputTag("trackExtenderWithMTD:generalTrackSigmaTofP"))
0124       ->setComment("Input ValueMap for track sigma(tof) as proton");
0125   desc.add<edm::InputTag>("vtxsSrc", edm::InputTag("unsortedOfflinePrimaryVertices4DwithPID"))
0126       ->setComment("Input primary vertex collection");
0127   desc.add<edm::InputTag>("trackMTDTimeQualityVMapTag", edm::InputTag("mtdTrackQualityMVA:mtdQualMVA"))
0128       ->setComment("Track MVA quality value");
0129   desc.add<double>("vtxMaxSigmaT", 0.025)
0130       ->setComment("Maximum primary vertex time uncertainty for use in particle id [ns]");
0131   desc.add<double>("maxDz", 0.1)
0132       ->setComment("Maximum distance in z for track-primary vertex association for particle id [cm]");
0133   desc.add<double>("maxDtSignificance", 5.0)
0134       ->setComment(
0135           "Maximum distance in time (normalized by uncertainty) for track-primary vertex association for particle id");
0136   desc.add<double>("minProbHeavy", 0.75)
0137       ->setComment("Minimum probability for a particle to be a kaon or proton before reassigning the timestamp");
0138   desc.add<double>("fixedT0Error", 0.)->setComment("Use a fixed T0 uncertainty [ns]");
0139   desc.add<double>("probPion", 1.)->setComment("A priori probability pions");
0140   desc.add<double>("probKaon", 1.)->setComment("A priori probability kaons");
0141   desc.add<double>("probProton", 1.)->setComment("A priori probability for protons");
0142   desc.add<double>("minTrackTimeQuality", 0.8)->setComment("Minimum MVA Quality selection on tracks");
0143   desc.add<bool>("MVASel", false)->setComment("Use MVA Quality selection");
0144   desc.add<bool>("vertexReassignment", true)->setComment("Track-vertex reassignment");
0145 
0146   descriptions.add("tofPIDProducer", desc);
0147 }
0148 
0149 template <class H, class T>
0150 void TOFPIDProducer::fillValueMap(edm::Event& iEvent,
0151                                   const edm::Handle<H>& handle,
0152                                   const std::vector<T>& vec,
0153                                   const std::string& name) const {
0154   auto out = std::make_unique<edm::ValueMap<T>>();
0155   typename edm::ValueMap<T>::Filler filler(*out);
0156   filler.insert(handle, vec.begin(), vec.end());
0157   filler.fill();
0158   iEvent.put(std::move(out), name);
0159 }
0160 
0161 void TOFPIDProducer::produce(edm::Event& ev, const edm::EventSetup& es) {
0162   edm::Handle<reco::TrackCollection> tracksH;
0163   ev.getByToken(tracksToken_, tracksH);
0164   const auto& tracks = *tracksH;
0165 
0166   const auto& t0In = ev.get(t0Token_);
0167 
0168   const auto& tmtdIn = ev.get(tmtdToken_);
0169 
0170   const auto& sigmat0In = ev.get(sigmat0Token_);
0171 
0172   const auto& sigmatmtdIn = ev.get(sigmatmtdToken_);
0173 
0174   const auto& tofkIn = ev.get(tofkToken_);
0175 
0176   const auto& tofpIn = ev.get(tofpToken_);
0177 
0178   const auto& sigmatofpiIn = ev.get(sigmatofpiToken_);
0179 
0180   const auto& sigmatofkIn = ev.get(sigmatofkToken_);
0181 
0182   const auto& sigmatofpIn = ev.get(sigmatofpToken_);
0183 
0184   const auto& vtxs = ev.get(vtxsToken_);
0185 
0186   const auto& trackMVAQualIn = ev.get(trackMTDTimeQualityToken_);
0187 
0188   //output value maps (PID probabilities and recalculated time at beamline)
0189   std::vector<float> t0OutRaw;
0190   std::vector<float> sigmat0OutRaw;
0191   std::vector<float> t0safeOutRaw;
0192   std::vector<float> sigmat0safeOutRaw;
0193   std::vector<float> probPiOutRaw;
0194   std::vector<float> probKOutRaw;
0195   std::vector<float> probPOutRaw;
0196 
0197   //Do work here
0198   for (unsigned int itrack = 0; itrack < tracks.size(); ++itrack) {
0199     const reco::Track& track = tracks[itrack];
0200     const reco::TrackRef trackref(tracksH, itrack);
0201     float t0 = t0In[trackref];
0202     float t0safe = t0;
0203     float sigmat0safe = sigmat0In[trackref];
0204     float sigmatmtd = (sigmatmtdIn[trackref] > 0. && fixedT0Error_ > 0.) ? fixedT0Error_ : sigmatmtdIn[trackref];
0205     float sigmat0 = sigmatmtd;
0206     float sigmatofpi = sigmatofpiIn[trackref];
0207     float sigmatofk = sigmatofkIn[trackref];
0208     float sigmatofp = sigmatofpIn[trackref];
0209 
0210     float prob_pi = -1.;
0211     float prob_k = -1.;
0212     float prob_p = -1.;
0213 
0214     float trackMVAQual = trackMVAQualIn[trackref];
0215 
0216     if (sigmat0 > 0. && (!MVASel_ || (MVASel_ && trackMVAQual >= minTrackTimeQuality_))) {
0217       double rsigmazsq = 1. / track.dzError() / track.dzError();
0218       std::array<double, 3> rsigmat = {{1. / std::sqrt(sigmatmtd * sigmatmtd + sigmatofpi * sigmatofpi),
0219                                         1. / std::sqrt(sigmatmtd * sigmatmtd + sigmatofk * sigmatofk),
0220                                         1. / std::sqrt(sigmatmtd * sigmatmtd + sigmatofp * sigmatofp)}};
0221 
0222       //find associated vertex
0223       int vtxidx = -1;
0224       int vtxidxmindz = -1;
0225       int vtxidxminchisq = -1;
0226       double mindz = maxDz_;
0227       double minchisq = std::numeric_limits<double>::max();
0228       //first try based on association weights, but keep track of closest in z and z-t as well
0229       for (unsigned int ivtx = 0; ivtx < vtxs.size(); ++ivtx) {
0230         const reco::Vertex& vtx = vtxs[ivtx];
0231         float w = vtx.trackWeight(trackref);
0232         if (w > 0.5) {
0233           vtxidx = ivtx;
0234           break;
0235         }
0236         double dz = std::abs(track.dz(vtx.position()));
0237         if (dz < mindz) {
0238           mindz = dz;
0239           vtxidxmindz = ivtx;
0240         }
0241         if (vtx.tError() > 0. && vtx.tError() < vtxMaxSigmaT_) {
0242           double dt = std::abs(t0 - vtx.t());
0243           double dtsig = dt * rsigmat[0];  //pion hp. uncertainty
0244           double chisq = dz * dz * rsigmazsq + dtsig * dtsig;
0245           if (dz < maxDz_ && dtsig < maxDtSignificance_ && chisq < minchisq) {
0246             minchisq = chisq;
0247             vtxidxminchisq = ivtx;
0248           }
0249         }
0250       }
0251 
0252       //if no vertex found based on association weights, fall back to closest in z or z-t
0253       if (vtxidx < 0) {
0254         //if closest vertex in z does not have valid time information, just use it,
0255         //otherwise use the closest vertex in z-t plane with timing info, with a fallback to the closest in z
0256         if (vtxidxmindz >= 0 && !(vtxs[vtxidxmindz].tError() > 0. && vtxs[vtxidxmindz].tError() < vtxMaxSigmaT_)) {
0257           vtxidx = vtxidxmindz;
0258         } else if (vtxidxminchisq >= 0) {
0259           vtxidx = vtxidxminchisq;
0260         } else if (vtxidxmindz >= 0) {
0261           vtxidx = vtxidxmindz;
0262         }
0263       }
0264 
0265       //testing mass hypotheses only possible if there is an associated vertex with time information
0266       if (vtxidx >= 0 && vtxs[vtxidx].tError() > 0. && vtxs[vtxidx].tError() < vtxMaxSigmaT_) {
0267         //compute chisq in z-t plane for nominal vertex and mass hypothesis (pion)
0268         const reco::Vertex& vtxnom = vtxs[vtxidx];
0269         double dznom = std::abs(track.dz(vtxnom.position()));
0270         double dtnom = std::abs(t0 - vtxnom.t());
0271         double dtsignom = dtnom * rsigmat[0];
0272         double chisqnom = dznom * dznom * rsigmazsq + dtsignom * dtsignom;
0273 
0274         //recompute t0 for alternate mass hypotheses
0275         double t0_best = t0;
0276 
0277         //reliable match, revert to raw mtd time uncertainty + tof uncertainty for pion hp
0278         if (dtsignom < maxDtSignificance_) {
0279           sigmat0safe = 1. / rsigmat[0];
0280           sigmat0 = sigmat0safe;
0281         }
0282 
0283         double tmtd = tmtdIn[trackref];
0284         double t0_k = tmtd - tofkIn[trackref];
0285         double t0_p = tmtd - tofpIn[trackref];
0286 
0287         double chisqmin = chisqnom;
0288 
0289         double chisqmin_pi = chisqnom;
0290         double chisqmin_k = std::numeric_limits<double>::max();
0291         double chisqmin_p = std::numeric_limits<double>::max();
0292         //loop through vertices and check for better matches
0293         for (unsigned int ivtx = 0; ivtx < vtxs.size(); ++ivtx) {
0294           const reco::Vertex& vtx = vtxs[ivtx];
0295           if (!vertexReassignment_) {
0296             if (ivtx != (unsigned int)vtxidx)
0297               continue;
0298           }
0299           if (!(vtx.tError() > 0. && vtx.tError() < vtxMaxSigmaT_)) {
0300             continue;
0301           }
0302 
0303           double dz = std::abs(track.dz(vtx.position()));
0304           if (dz >= maxDz_) {
0305             continue;
0306           }
0307 
0308           double chisqdz = dz * dz * rsigmazsq;
0309 
0310           double dt_k = std::abs(t0_k - vtx.t());
0311           double dtsig_k = dt_k * rsigmat[1];
0312           double chisq_k = chisqdz + dtsig_k * dtsig_k;
0313 
0314           if (dtsig_k < maxDtSignificance_ && chisq_k < chisqmin_k) {
0315             chisqmin_k = chisq_k;
0316           }
0317 
0318           double dt_p = std::abs(t0_p - vtx.t());
0319           double dtsig_p = dt_p * rsigmat[2];
0320           double chisq_p = chisqdz + dtsig_p * dtsig_p;
0321 
0322           if (dtsig_p < maxDtSignificance_ && chisq_p < chisqmin_p) {
0323             chisqmin_p = chisq_p;
0324           }
0325 
0326           if (dtsig_k < maxDtSignificance_ && chisq_k < chisqmin) {
0327             chisqmin = chisq_k;
0328             t0_best = t0_k;
0329             t0safe = t0_k;
0330             sigmat0safe = 1. / rsigmat[1];
0331             sigmat0 = sigmat0safe;
0332           }
0333           if (dtsig_p < maxDtSignificance_ && chisq_p < chisqmin) {
0334             chisqmin = chisq_p;
0335             t0_best = t0_p;
0336             t0safe = t0_p;
0337             sigmat0safe = 1. / rsigmat[2];
0338             sigmat0 = sigmat0safe;
0339           }
0340         }
0341 
0342         //compute PID probabilities
0343         //*TODO* deal with heavier nucleons and/or BSM case here?
0344         double rawprob_pi = probPion_ * exp(-0.5 * chisqmin_pi);
0345         double rawprob_k = probKaon_ * exp(-0.5 * chisqmin_k);
0346         double rawprob_p = probProton_ * exp(-0.5 * chisqmin_p);
0347 
0348         double normprob = 1. / (rawprob_pi + rawprob_k + rawprob_p);
0349 
0350         prob_pi = rawprob_pi * normprob;
0351         prob_k = rawprob_k * normprob;
0352         prob_p = rawprob_p * normprob;
0353 
0354         double prob_heavy = 1. - prob_pi;
0355 
0356         if (prob_heavy > minProbHeavy_) {
0357           t0 = t0_best;
0358         }
0359       }
0360     }
0361 
0362     t0OutRaw.push_back(t0);
0363     sigmat0OutRaw.push_back(sigmat0);
0364     t0safeOutRaw.push_back(t0safe);
0365     sigmat0safeOutRaw.push_back(sigmat0safe);
0366     probPiOutRaw.push_back(prob_pi);
0367     probKOutRaw.push_back(prob_k);
0368     probPOutRaw.push_back(prob_p);
0369   }
0370 
0371   fillValueMap(ev, tracksH, t0OutRaw, t0Name);
0372   fillValueMap(ev, tracksH, sigmat0OutRaw, sigmat0Name);
0373   fillValueMap(ev, tracksH, t0safeOutRaw, t0safeName);
0374   fillValueMap(ev, tracksH, sigmat0safeOutRaw, sigmat0safeName);
0375   fillValueMap(ev, tracksH, probPiOutRaw, probPiName);
0376   fillValueMap(ev, tracksH, probKOutRaw, probKName);
0377   fillValueMap(ev, tracksH, probPOutRaw, probPName);
0378 }
0379 
0380 //define this as a plug-in
0381 #include <FWCore/Framework/interface/MakerMacros.h>
0382 DEFINE_FWK_MODULE(TOFPIDProducer);
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.2.1 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!