File indexing completed on 2024-04-25 02:14:10
0001 #include <vector>
0002 #include <limits>
0003 #include <string>
0004 #include <atomic>
0005
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "FWCore/Framework/interface/EventSetup.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009 #include "FWCore/Framework/interface/global/EDProducer.h"
0010 #include "FWCore/ParameterSet/interface/FileInPath.h"
0011 #include "FWCore/Utilities/interface/InputTag.h"
0012 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0013 #include "FWCore/Utilities/interface/EDGetToken.h"
0014 #include "FWCore/Utilities/interface/Exception.h"
0015 #include "FWCore/Utilities/interface/InputTag.h"
0016
0017 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0018 #include "DataFormats/Candidate/interface/VertexCompositeCandidate.h"
0019 #include "DataFormats/Common/interface/Handle.h"
0020 #include "DataFormats/Common/interface/View.h"
0021 #include "DataFormats/Math/interface/Vector3D.h"
0022 #include "DataFormats/VertexReco/interface/Vertex.h"
0023
0024 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0025
0026 #include "RecoTracker/DisplacedRegionalTracking/plugins/DisplacedVertexCluster.h"
0027
0028 using namespace std;
0029 typedef DisplacedVertexCluster::Distance Distance;
0030 typedef DisplacedVertexCluster::DistanceItr DistanceItr;
0031
0032 class DisplacedRegionSeedingVertexProducer : public edm::global::EDProducer<> {
0033 public:
0034 DisplacedRegionSeedingVertexProducer(const edm::ParameterSet &);
0035 ~DisplacedRegionSeedingVertexProducer() override;
0036 void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
0037 static void fillDescriptions(edm::ConfigurationDescriptions &);
0038
0039 private:
0040
0041 const double rParam_;
0042
0043
0044 const double minRadius_;
0045 const double nearThreshold_;
0046 const double farThreshold_;
0047 const double discriminatorCut_;
0048 const unsigned int maxPseudoROIs_;
0049 const vector<string> input_names_;
0050 const vector<string> output_names_;
0051
0052 const edm::EDGetTokenT<reco::BeamSpot> beamSpotToken_;
0053 const edm::EDGetTokenT<edm::View<reco::VertexCompositeCandidate> > trackClustersToken_;
0054
0055 tensorflow::Session *session_;
0056
0057 double getDiscriminatorValue(const DisplacedVertexCluster &, const reco::BeamSpot &) const;
0058 };
0059
0060 DisplacedRegionSeedingVertexProducer::DisplacedRegionSeedingVertexProducer(const edm::ParameterSet &cfg)
0061 : rParam_(cfg.getParameter<double>("rParam")),
0062 minRadius_(cfg.getParameter<double>("minRadius")),
0063 nearThreshold_(cfg.getParameter<double>("nearThreshold")),
0064 farThreshold_(cfg.getParameter<double>("farThreshold")),
0065 discriminatorCut_(cfg.getParameter<double>("discriminatorCut")),
0066 maxPseudoROIs_(cfg.getParameter<unsigned int>("maxPseudoROIs")),
0067 input_names_(cfg.getParameter<vector<string> >("input_names")),
0068 output_names_(cfg.getParameter<vector<string> >("output_names")),
0069 beamSpotToken_(consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))),
0070 trackClustersToken_(
0071 consumes<edm::View<reco::VertexCompositeCandidate> >(cfg.getParameter<edm::InputTag>("trackClusters"))) {
0072 unsigned nThreads = cfg.getUntrackedParameter<unsigned>("nThreads");
0073 tensorflow::Options options;
0074 options.setThreading(nThreads);
0075 string pbFile = cfg.getParameter<edm::FileInPath>("graph_path").fullPath();
0076 auto graphDef = tensorflow::loadGraphDef(pbFile);
0077 session_ = tensorflow::createSession(graphDef, options);
0078
0079 produces<vector<reco::Vertex> >("nearRegionsOfInterest");
0080 produces<vector<reco::Vertex> >("farRegionsOfInterest");
0081 }
0082
0083 DisplacedRegionSeedingVertexProducer::~DisplacedRegionSeedingVertexProducer() {
0084 if (session_ != nullptr)
0085 tensorflow::closeSession(session_);
0086 }
0087
0088 void DisplacedRegionSeedingVertexProducer::produce(edm::StreamID streamID,
0089 edm::Event &event,
0090 const edm::EventSetup &setup) const {
0091 const auto &beamSpot = event.get(beamSpotToken_);
0092 const math::XYZVector bs(beamSpot.position());
0093
0094 const auto &trackClusters = event.get(trackClustersToken_);
0095
0096
0097 vector<DisplacedVertexCluster> pseudoROIs;
0098 pseudoROIs.reserve(std::min(maxPseudoROIs_, trackClusters.size()));
0099 vector<Distance> distances;
0100 const double minTrackClusterRadius = minRadius_ - rParam_;
0101 for (unsigned i = 0; i < trackClusters.size(); i++) {
0102 const reco::VertexCompositeCandidate &trackCluster = trackClusters[i];
0103 const math::XYZVector x(trackCluster.vertex());
0104 if (minRadius_ < 0.0 || minTrackClusterRadius < 0.0 || (x - bs).rho() > minTrackClusterRadius)
0105 pseudoROIs.emplace_back(&trackClusters.at(i), rParam_);
0106 if (pseudoROIs.size() == maxPseudoROIs_) {
0107 edm::LogWarning("DisplacedRegionSeedingVertexProducer")
0108 << "Truncated list of pseudoROIs at " << maxPseudoROIs_ << " out of " << trackClusters.size()
0109 << " possible track clusters.";
0110 break;
0111 }
0112 }
0113 if (pseudoROIs.size() > 1) {
0114 distances.reserve(pseudoROIs.size() * std::max(1.0, pseudoROIs.size() * 0.05));
0115 DisplacedVertexClusterItr secondToLast = pseudoROIs.end();
0116 secondToLast--;
0117 for (DisplacedVertexClusterItr i = pseudoROIs.begin(); i != secondToLast; i++) {
0118 DisplacedVertexClusterItr j = i;
0119 j++;
0120 for (; j != pseudoROIs.end(); j++) {
0121 distances.emplace_back(i, j);
0122
0123
0124
0125 if (distances.back().distance2() > 16.0 * rParam_ * rParam_)
0126 distances.pop_back();
0127 }
0128 }
0129 }
0130
0131 auto itBegin = distances.begin();
0132 auto itLast = distances.end();
0133
0134
0135 while (itBegin != itLast) {
0136
0137
0138
0139
0140 DistanceItr dBest = itBegin;
0141 double distanceBest = dBest->distance2();
0142
0143 for (auto i = itBegin; i != itLast; i++) {
0144 if (distanceBest > i->distance2()) {
0145 dBest = i;
0146 distanceBest = i->distance2();
0147 }
0148 }
0149
0150 if (dBest->distance2() > rParam_ * rParam_)
0151 break;
0152
0153 dBest->entities().first->merge(*dBest->entities().second);
0154 dBest->entities().second->setInvalid();
0155
0156 const auto distancePred = [](const Distance &a) {
0157 return (!a.entities().first->valid() || !a.entities().second->valid());
0158 };
0159 itLast = std::remove_if(itBegin, itLast, distancePred);
0160 }
0161
0162 const auto pseudoROIPred = [](const DisplacedVertexCluster &a) { return !a.valid(); };
0163 auto remove_invalid = std::remove_if(pseudoROIs.begin(), pseudoROIs.end(), pseudoROIPred);
0164
0165
0166 const auto roiPred = [&](const DisplacedVertexCluster &roi) {
0167 if (!roi.valid())
0168 return true;
0169 const auto &x(roi.centerOfMass());
0170 if ((x - bs).rho() < minRadius_)
0171 return true;
0172 const double discriminatorValue = ((discriminatorCut_ > 0.0) ? getDiscriminatorValue(roi, beamSpot) : 1.0);
0173 if (discriminatorValue < discriminatorCut_)
0174 return true;
0175 return false;
0176 };
0177 auto remove_pred = std::remove_if(pseudoROIs.begin(), remove_invalid, roiPred);
0178
0179 auto nearRegionsOfInterest = make_unique<vector<reco::Vertex> >();
0180 auto farRegionsOfInterest = make_unique<vector<reco::Vertex> >();
0181
0182 constexpr std::array<double, 6> errorA{{1.0, 0.0, 1.0, 0.0, 0.0, 1.0}};
0183 static const reco::Vertex::Error errorRegion(errorA.begin(), errorA.end(), true, true);
0184
0185 for (auto it = pseudoROIs.begin(); it != remove_pred; ++it) {
0186 auto const &roi = *it;
0187 const auto &x(roi.centerOfMass());
0188 if ((x - bs).rho() < nearThreshold_)
0189 nearRegionsOfInterest->emplace_back(reco::Vertex::Point(roi.centerOfMass()), errorRegion);
0190 if ((x - bs).rho() > farThreshold_)
0191 farRegionsOfInterest->emplace_back(reco::Vertex::Point(roi.centerOfMass()), errorRegion);
0192 }
0193
0194 event.put(std::move(nearRegionsOfInterest), "nearRegionsOfInterest");
0195 event.put(std::move(farRegionsOfInterest), "farRegionsOfInterest");
0196 }
0197
0198 void DisplacedRegionSeedingVertexProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0199 edm::ParameterSetDescription desc;
0200
0201 desc.add<double>("rParam", 1.0);
0202 desc.add<double>("minRadius", -1.0);
0203 desc.add<double>("nearThreshold", 9999.0);
0204 desc.add<double>("farThreshold", -1.0);
0205 desc.add<double>("discriminatorCut", -1.0);
0206 desc.add<vector<string> >("input_names", {"phi_0", "phi_1"});
0207 desc.add<vector<string> >("output_names", {"model_5/activation_10/Softmax"});
0208 desc.add<unsigned int>("maxPseudoROIs", 10000);
0209 desc.addUntracked<unsigned>("nThreads", 1);
0210 desc.add<edm::FileInPath>(
0211 "graph_path",
0212 edm::FileInPath(
0213 "RecoTracker/DisplacedRegionalTracking/data/FullData_Phi-64-128-256_16-32-64_F-128-64-32_Model.pb"));
0214 desc.add<edm::InputTag>("beamSpot", edm::InputTag("offlineBeamSpot"));
0215 desc.add<edm::InputTag>("trackClusters", edm::InputTag("generalV0Candidates", "Kshort"));
0216
0217 descriptions.add("displacedRegionProducer", desc);
0218 }
0219
0220 double DisplacedRegionSeedingVertexProducer::getDiscriminatorValue(const DisplacedVertexCluster &roi,
0221 const reco::BeamSpot &bs) const {
0222
0223
0224
0225
0226 constexpr int maxNVertices = 40;
0227 constexpr int nVertexFeatures = 23;
0228 constexpr int nDimVertices = 3;
0229
0230 constexpr int maxNAnnulusTracks = 10;
0231 constexpr int nAnnulusTrackFeatures = 8;
0232 constexpr int nDimAnnulusTracks = 3;
0233
0234 tensorflow::Tensor vertexTensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({1, maxNVertices, nVertexFeatures}));
0235 auto vertex_map = vertexTensor.tensor<float, nDimVertices>();
0236 tensorflow::Tensor annulusTensor(tensorflow::DT_FLOAT,
0237 tensorflow::TensorShape({1, maxNAnnulusTracks, nAnnulusTrackFeatures}));
0238 auto annulus_map = annulusTensor.tensor<float, nDimAnnulusTracks>();
0239
0240 for (int i = 0, map_i = 0; map_i < maxNVertices; i++, map_i++) {
0241 if (i >= static_cast<int>(roi.nConstituents()))
0242 for (unsigned j = 0; j < nVertexFeatures; j++)
0243 vertex_map(0, map_i, j) = 0.0;
0244 else {
0245 const auto &trackCluster = *roi.constituent(i);
0246 const auto &track0 = *trackCluster.daughter(0)->bestTrack();
0247 const auto &track1 = *trackCluster.daughter(1)->bestTrack();
0248
0249 vertex_map(0, map_i, 0) = trackCluster.vx() - bs.x0();
0250 vertex_map(0, map_i, 1) = trackCluster.vy() - bs.y0();
0251 vertex_map(0, map_i, 2) = trackCluster.vz() - bs.z0();
0252
0253 vertex_map(0, map_i, 3) = trackCluster.vertexCovariance()(0, 0);
0254 vertex_map(0, map_i, 4) = trackCluster.vertexCovariance()(0, 1);
0255 vertex_map(0, map_i, 5) = trackCluster.vertexCovariance()(0, 2);
0256 vertex_map(0, map_i, 6) = trackCluster.vertexCovariance()(1, 1);
0257 vertex_map(0, map_i, 7) = trackCluster.vertexCovariance()(1, 2);
0258 vertex_map(0, map_i, 8) = trackCluster.vertexCovariance()(2, 2);
0259
0260 vertex_map(0, map_i, 9) = track0.charge() * track0.pt();
0261 vertex_map(0, map_i, 10) = track0.eta();
0262 vertex_map(0, map_i, 11) = track0.phi();
0263 vertex_map(0, map_i, 12) = track0.dxy(bs);
0264 vertex_map(0, map_i, 13) = track0.dz(bs.position());
0265 vertex_map(0, map_i, 14) = track0.normalizedChi2();
0266 vertex_map(0, map_i, 15) = track0.quality(reco::Track::highPurity) ? 1 : 0;
0267
0268 vertex_map(0, map_i, 16) = track1.charge() * track1.pt();
0269 vertex_map(0, map_i, 17) = track1.eta();
0270 vertex_map(0, map_i, 18) = track1.phi();
0271 vertex_map(0, map_i, 19) = track1.dxy(bs);
0272 vertex_map(0, map_i, 20) = track1.dz(bs.position());
0273 vertex_map(0, map_i, 21) = track1.normalizedChi2();
0274 vertex_map(0, map_i, 22) = track1.quality(reco::Track::highPurity) ? 1 : 0;
0275 }
0276 }
0277
0278 for (int i = 0; i < maxNAnnulusTracks; i++)
0279 for (unsigned j = 0; j < nAnnulusTrackFeatures; j++)
0280 annulus_map(0, i, j) = 0.0;
0281
0282 tensorflow::NamedTensorList input_tensors;
0283 input_tensors.resize(2);
0284 input_tensors[0] = tensorflow::NamedTensor(input_names_.at(0), vertexTensor);
0285 input_tensors[1] = tensorflow::NamedTensor(input_names_.at(1), annulusTensor);
0286 vector<tensorflow::Tensor> outputs;
0287 tensorflow::run(session_, input_tensors, output_names_, &outputs);
0288
0289 return (outputs.at(0).flat<float>()(1));
0290 }
0291
0292 DEFINE_FWK_MODULE(DisplacedRegionSeedingVertexProducer);