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