File indexing completed on 2023-03-17 11:20:18
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
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
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
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
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
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
0203 if (vtxidx < 0) {
0204
0205
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
0216 if (vtxidx >= 0 && vtxs[vtxidx].tError() > 0. && vtxs[vtxidx].tError() < vtxMaxSigmaT_) {
0217
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
0225 double t0_best = t0;
0226
0227
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
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
0285
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
0323 #include <FWCore/Framework/interface/MakerMacros.h>
0324 DEFINE_FWK_MODULE(TOFPIDProducer);