DisplacedRegionalTracking/plugins/DisplacedRegionSeedingVertexProducer.cc

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   // clustering parameters
0041   const double rParam_;
0042
0043   // selection parameters
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   // Initialize distances.
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         // Track clusters farther apart than 4 times rParam_ (i.e., 16 times
0124         // rParam_^2) cannot wind up in the same ROI, so remove these pairs.
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   // Do clustering.
0135   while (itBegin != itLast) {
0136     //find the lowest distance. Lots of repeatitive calculations done here
0137     //as from loop iteration to loop iteration only sqrt(distances.size()) distances
0138     //need to be recomputed (those involving best_i
0139     //but this is much better than sorting distances..
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   // Remove invalid ROIs.
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   // The network takes in two maps of data features, one with information
0223   // related to the pairwise track vertices and one with information related to
0224   // the tracks in an isolation annulus.
0225
0226   constexpr int maxNVertices = 40;     // maximum number of pairwise track vertices per ROI
0227   constexpr int nVertexFeatures = 23;  // number of features per pairwise track vertex
0228   constexpr int nDimVertices = 3;      // number of tensor dimensions for the map of pairwise track vertices
0229
0230   constexpr int maxNAnnulusTracks = 10;     // maximum number of annulus tracks per ROI
0231   constexpr int nAnnulusTrackFeatures = 8;  // number of features per annulus track
0232   constexpr int nDimAnnulusTracks = 3;      // number of tensor dimensions for the map of annulus tracks
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);