DisplacedRegionalTracking/plugins/DisplacedRegionSeedingVertexProducer.cc

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   // clustering parameters
0042   const double rParam_;
0043
0044   // selection parameters
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   // Initialize distances.
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         // Track clusters farther apart than 4 times rParam_ (i.e., 16 times
0115         // rParam_^2) cannot wind up in the same ROI, so remove these pairs.
0116         if (distances.back().distance2() > 16.0 * rParam_ * rParam_)
0117           distances.pop_back();
0118       }
0119     }
0120   }
0121
0122   // Do clustering.
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   // Remove invalid ROIs.
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   // The network takes in two maps of data features, one with information
0197   // related to the pairwise track vertices and one with information related to
0198   // the tracks in an isolation annulus.
0199
0200   constexpr int maxNVertices = 40;     // maximum number of pairwise track vertices per ROI
0201   constexpr int nVertexFeatures = 23;  // number of features per pairwise track vertex
0202   constexpr int nDimVertices = 3;      // number of tensor dimensions for the map of pairwise track vertices
0203
0204   constexpr int maxNAnnulusTracks = 10;     // maximum number of annulus tracks per ROI
0205   constexpr int nAnnulusTrackFeatures = 8;  // number of features per annulus track
0206   constexpr int nDimAnnulusTracks = 3;      // number of tensor dimensions for the map of annulus tracks
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);