TimingIDTools/plugins/TOFPIDProducer.cc

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/GlobalSystemOfUnits.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<reco::VertexCollection> vtxsToken_;
0052   double vtxMaxSigmaT_;
0053   double maxDz_;
0054   double maxDtSignificance_;
0055   double minProbHeavy_;
0056   double fixedT0Error_;
0057 };
0058
0059 TOFPIDProducer::TOFPIDProducer(const ParameterSet& iConfig)
0060     : tracksToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("tracksSrc"))),
0061       t0Token_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("t0Src"))),
0062       tmtdToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("tmtdSrc"))),
0063       sigmat0Token_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmat0Src"))),
0064       sigmatmtdToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("sigmatmtdSrc"))),
0065       tofkToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("tofkSrc"))),
0066       tofpToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("tofpSrc"))),
0067       vtxsToken_(consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("vtxsSrc"))),
0068       vtxMaxSigmaT_(iConfig.getParameter<double>("vtxMaxSigmaT")),
0069       maxDz_(iConfig.getParameter<double>("maxDz")),
0070       maxDtSignificance_(iConfig.getParameter<double>("maxDtSignificance")),
0071       minProbHeavy_(iConfig.getParameter<double>("minProbHeavy")),
0072       fixedT0Error_(iConfig.getParameter<double>("fixedT0Error")) {
0073   produces<edm::ValueMap<float>>(t0Name);
0074   produces<edm::ValueMap<float>>(sigmat0Name);
0075   produces<edm::ValueMap<float>>(t0safeName);
0076   produces<edm::ValueMap<float>>(sigmat0safeName);
0077   produces<edm::ValueMap<float>>(probPiName);
0078   produces<edm::ValueMap<float>>(probKName);
0079   produces<edm::ValueMap<float>>(probPName);
0080 }
0081
0082 // Configuration descriptions
0083 void TOFPIDProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0084   edm::ParameterSetDescription desc;
0085   desc.add<edm::InputTag>("tracksSrc", edm::InputTag("generalTracks"))->setComment("Input tracks collection");
0086   desc.add<edm::InputTag>("t0Src", edm::InputTag("trackExtenderWithMTD:generalTrackt0"))
0087       ->setComment("Input ValueMap for track time at beamline");
0088   desc.add<edm::InputTag>("tmtdSrc", edm::InputTag("trackExtenderWithMTD:generalTracktmtd"))
0089       ->setComment("Input ValueMap for track time at MTD");
0090   desc.add<edm::InputTag>("sigmat0Src", edm::InputTag("trackExtenderWithMTD:generalTracksigmat0"))
0091       ->setComment("Input ValueMap for track time uncertainty at beamline");
0092   desc.add<edm::InputTag>("sigmatmtdSrc", edm::InputTag("trackExtenderWithMTD:generalTracksigmatmtd"))
0093       ->setComment("Input ValueMap for track time uncertainty at MTD");
0094   desc.add<edm::InputTag>("tofkSrc", edm::InputTag("trackExtenderWithMTD:generalTrackTofK"))
0095       ->setComment("Input ValueMap for track tof as kaon");
0096   desc.add<edm::InputTag>("tofpSrc", edm::InputTag("trackExtenderWithMTD:generalTrackTofP"))
0097       ->setComment("Input ValueMap for track tof as proton");
0098   desc.add<edm::InputTag>("vtxsSrc", edm::InputTag("unsortedOfflinePrimaryVertices4DwithPID"))
0099       ->setComment("Input primary vertex collection");
0100   desc.add<double>("vtxMaxSigmaT", 0.025)
0101       ->setComment("Maximum primary vertex time uncertainty for use in particle id [ns]");
0102   desc.add<double>("maxDz", 0.1)
0103       ->setComment("Maximum distance in z for track-primary vertex association for particle id [cm]");
0104   desc.add<double>("maxDtSignificance", 5.0)
0105       ->setComment(
0106           "Maximum distance in time (normalized by uncertainty) for track-primary vertex association for particle id");
0107   desc.add<double>("minProbHeavy", 0.75)
0108       ->setComment("Minimum probability for a particle to be a kaon or proton before reassigning the timestamp");
0109   desc.add<double>("fixedT0Error", 0.)->setComment("Use a fixed T0 uncertainty [ns]");
0110
0111   descriptions.add("tofPIDProducer", desc);
0112 }
0113
0114 template <class H, class T>
0115 void TOFPIDProducer::fillValueMap(edm::Event& iEvent,
0116                                   const edm::Handle<H>& handle,
0117                                   const std::vector<T>& vec,
0118                                   const std::string& name) const {
0119   auto out = std::make_unique<edm::ValueMap<T>>();
0120   typename edm::ValueMap<T>::Filler filler(*out);
0121   filler.insert(handle, vec.begin(), vec.end());
0122   filler.fill();
0123   iEvent.put(std::move(out), name);
0124 }
0125
0126 void TOFPIDProducer::produce(edm::Event& ev, const edm::EventSetup& es) {
0127   edm::Handle<reco::TrackCollection> tracksH;
0128   ev.getByToken(tracksToken_, tracksH);
0129   const auto& tracks = *tracksH;
0130
0131   const auto& t0In = ev.get(t0Token_);
0132
0133   const auto& tmtdIn = ev.get(tmtdToken_);
0134
0135   const auto& sigmat0In = ev.get(sigmat0Token_);
0136
0137   const auto& sigmatmtdIn = ev.get(sigmatmtdToken_);
0138
0139   const auto& tofkIn = ev.get(tofkToken_);
0140
0141   const auto& tofpIn = ev.get(tofpToken_);
0142
0143   const auto& vtxs = ev.get(vtxsToken_);
0144
0145   //output value maps (PID probabilities and recalculated time at beamline)
0146   std::vector<float> t0OutRaw;
0147   std::vector<float> sigmat0OutRaw;
0148   std::vector<float> t0safeOutRaw;
0149   std::vector<float> sigmat0safeOutRaw;
0150   std::vector<float> probPiOutRaw;
0151   std::vector<float> probKOutRaw;
0152   std::vector<float> probPOutRaw;
0153
0154   //Do work here
0155   for (unsigned int itrack = 0; itrack < tracks.size(); ++itrack) {
0156     const reco::Track& track = tracks[itrack];
0157     const reco::TrackRef trackref(tracksH, itrack);
0158     float t0 = t0In[trackref];
0159     float t0safe = t0;
0160     float sigmat0safe = sigmat0In[trackref];
0161     float sigmatmtd = (sigmatmtdIn[trackref] > 0. && fixedT0Error_ > 0.) ? fixedT0Error_ : sigmatmtdIn[trackref];
0162     float sigmat0 = sigmatmtd;
0163
0164     float prob_pi = -1.;
0165     float prob_k = -1.;
0166     float prob_p = -1.;
0167
0168     if (sigmat0 > 0.) {
0169       double rsigmazsq = 1. / track.dzError() / track.dzError();
0170       double rsigmat = 1. / sigmatmtd;
0171
0172       //find associated vertex
0173       int vtxidx = -1;
0174       int vtxidxmindz = -1;
0175       int vtxidxminchisq = -1;
0176       double mindz = maxDz_;
0177       double minchisq = std::numeric_limits<double>::max();
0178       //first try based on association weights, but keep track of closest in z and z-t as well
0179       for (unsigned int ivtx = 0; ivtx < vtxs.size(); ++ivtx) {
0180         const reco::Vertex& vtx = vtxs[ivtx];
0181         float w = vtx.trackWeight(trackref);
0182         if (w > 0.5) {
0183           vtxidx = ivtx;
0184           break;
0185         }
0186         double dz = std::abs(track.dz(vtx.position()));
0187         if (dz < mindz) {
0188           mindz = dz;
0189           vtxidxmindz = ivtx;
0190         }
0191         if (vtx.tError() > 0. && vtx.tError() < vtxMaxSigmaT_) {
0192           double dt = std::abs(t0 - vtx.t());
0193           double dtsig = dt * rsigmat;
0194           double chisq = dz * dz * rsigmazsq + dtsig * dtsig;
0195           if (dz < maxDz_ && dtsig < maxDtSignificance_ && chisq < minchisq) {
0196             minchisq = chisq;
0197             vtxidxminchisq = ivtx;
0198           }
0199         }
0200       }
0201
0202       //if no vertex found based on association weights, fall back to closest in z or z-t
0203       if (vtxidx < 0) {
0204         //if closest vertex in z does not have valid time information, just use it,
0205         //otherwise use the closest vertex in z-t plane with timing info, with a fallback to the closest in z
0206         if (vtxidxmindz >= 0 && !(vtxs[vtxidxmindz].tError() > 0. && vtxs[vtxidxmindz].tError() < vtxMaxSigmaT_)) {
0207           vtxidx = vtxidxmindz;
0208         } else if (vtxidxminchisq >= 0) {
0209           vtxidx = vtxidxminchisq;
0210         } else if (vtxidxmindz >= 0) {
0211           vtxidx = vtxidxmindz;
0212         }
0213       }
0214
0215       //testing mass hypotheses only possible if there is an associated vertex with time information
0216       if (vtxidx >= 0 && vtxs[vtxidx].tError() > 0. && vtxs[vtxidx].tError() < vtxMaxSigmaT_) {
0217         //compute chisq in z-t plane for nominal vertex and mass hypothesis (pion)
0218         const reco::Vertex& vtxnom = vtxs[vtxidx];
0219         double dznom = std::abs(track.dz(vtxnom.position()));
0220         double dtnom = std::abs(t0 - vtxnom.t());
0221         double dtsignom = dtnom * rsigmat;
0222         double chisqnom = dznom * dznom * rsigmazsq + dtsignom * dtsignom;
0223
0224         //recompute t0 for alternate mass hypotheses
0225         double t0_best = t0;
0226
0227         //reliable match, revert to raw mtd time uncertainty
0228         if (dtsignom < maxDtSignificance_) {
0229           sigmat0safe = sigmatmtd;
0230         }
0231
0232         double tmtd = tmtdIn[trackref];
0233         double t0_k = tmtd - tofkIn[trackref];
0234         double t0_p = tmtd - tofpIn[trackref];
0235
0236         double chisqmin = chisqnom;
0237
0238         double chisqmin_pi = chisqnom;
0239         double chisqmin_k = std::numeric_limits<double>::max();
0240         double chisqmin_p = std::numeric_limits<double>::max();
0241         //loop through vertices and check for better matches
0242         for (const reco::Vertex& vtx : vtxs) {
0243           if (!(vtx.tError() > 0. && vtx.tError() < vtxMaxSigmaT_)) {
0244             continue;
0245           }
0246
0247           double dz = std::abs(track.dz(vtx.position()));
0248           if (dz >= maxDz_) {
0249             continue;
0250           }
0251
0252           double chisqdz = dz * dz * rsigmazsq;
0253
0254           double dt_k = std::abs(t0_k - vtx.t());
0255           double dtsig_k = dt_k * rsigmat;
0256           double chisq_k = chisqdz + dtsig_k * dtsig_k;
0257
0258           if (dtsig_k < maxDtSignificance_ && chisq_k < chisqmin_k) {
0259             chisqmin_k = chisq_k;
0260           }
0261
0262           double dt_p = std::abs(t0_p - vtx.t());
0263           double dtsig_p = dt_p * rsigmat;
0264           double chisq_p = chisqdz + dtsig_p * dtsig_p;
0265
0266           if (dtsig_p < maxDtSignificance_ && chisq_p < chisqmin_p) {
0267             chisqmin_p = chisq_p;
0268           }
0269
0270           if (dtsig_k < maxDtSignificance_ && chisq_k < chisqmin) {
0271             chisqmin = chisq_k;
0272             t0_best = t0_k;
0273             t0safe = t0_k;
0274             sigmat0safe = sigmatmtd;
0275           }
0276           if (dtsig_p < maxDtSignificance_ && chisq_p < chisqmin) {
0277             chisqmin = chisq_p;
0278             t0_best = t0_p;
0279             t0safe = t0_p;
0280             sigmat0safe = sigmatmtd;
0281           }
0282         }
0283
0284         //compute PID probabilities
0285         //*TODO* deal with heavier nucleons and/or BSM case here?
0286         double rawprob_pi = exp(-0.5 * chisqmin_pi);
0287         double rawprob_k = exp(-0.5 * chisqmin_k);
0288         double rawprob_p = exp(-0.5 * chisqmin_p);
0289
0290         double normprob = 1. / (rawprob_pi + rawprob_k + rawprob_p);
0291
0292         prob_pi = rawprob_pi * normprob;
0293         prob_k = rawprob_k * normprob;
0294         prob_p = rawprob_p * normprob;
0295
0296         double prob_heavy = 1. - prob_pi;
0297
0298         if (prob_heavy > minProbHeavy_) {
0299           t0 = t0_best;
0300         }
0301       }
0302     }
0303
0304     t0OutRaw.push_back(t0);
0305     sigmat0OutRaw.push_back(sigmat0);
0306     t0safeOutRaw.push_back(t0safe);
0307     sigmat0safeOutRaw.push_back(sigmat0safe);
0308     probPiOutRaw.push_back(prob_pi);
0309     probKOutRaw.push_back(prob_k);
0310     probPOutRaw.push_back(prob_p);
0311   }
0312
0313   fillValueMap(ev, tracksH, t0OutRaw, t0Name);
0314   fillValueMap(ev, tracksH, sigmat0OutRaw, sigmat0Name);
0315   fillValueMap(ev, tracksH, t0safeOutRaw, t0safeName);
0316   fillValueMap(ev, tracksH, sigmat0safeOutRaw, sigmat0safeName);
0317   fillValueMap(ev, tracksH, probPiOutRaw, probPiName);
0318   fillValueMap(ev, tracksH, probKOutRaw, probKName);
0319   fillValueMap(ev, tracksH, probPOutRaw, probPName);
0320 }
0321
0322 //define this as a plug-in
0323 #include <FWCore/Framework/interface/MakerMacros.h>
0324 DEFINE_FWK_MODULE(TOFPIDProducer);