Track classifier using TensorFlow

DataFormats/TrackTfGraph/BuildFile.xml

@@ -0,0 +1,5 @@
+<use name="PhysicsTools/TensorFlow"/>
+<use name="RecoTracker/Record"/>
+<export>
+  <lib name="1"/>
+</export>

DataFormats/TrackTfGraph/interface/TfGraphDefWrapper.h

@@ -0,0 +1,16 @@
+#ifndef TrackTfGraph_TfGraphDefWrapper_h
+#define TrackTfGraph_TfGraphDefWrapper_h
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+class TfGraphDefWrapper {
+    public:
+        TfGraphDefWrapper(tensorflow::GraphDef*);
+        tensorflow::GraphDef* GetGraphDef() const;
+
+    private:
+        tensorflow::GraphDef* graphDef_;
+};
+
+
+#endif

DataFormats/TrackTfGraph/plugins/BuildFile.xml

@@ -0,0 +1,11 @@
+<use name="FWCore/Framework"/>
+<use name="FWCore/ParameterSet"/>
+<use name="FWCore/PluginManager"/>
+<use name="PhysicsTools/TensorFlow"/>
+<use name="FWCore/Utilities"/>
+<use name="RecoTracker/Record"/>
+<use name="DataFormats/TrackTfGraph"/>
+<flags EDM_PLUGIN="1"/>
+<export>
+  <lib name="1"/>
+</export>

DataFormats/TrackTfGraph/plugins/TfGraphDefProducer.cc

@@ -0,0 +1,66 @@
+// -*- C++ -*-
+//
+// Package:    test/TFGraphDefProducer
+// Class:      TFGraphDefProducer
+//
+/**\class TFGraphDefProducer
+ Description: Produces TfGraphRecord into the event containing a tensorflow GraphDef object that can be used for running inference on a pretrained network
+*/
+//
+// Original Author:  Joona Havukainen
+//         Created:  Fri, 24 Jul 2020 08:04:00 GMT
+//
+//
+
+// system include files
+#include <memory>
+
+// user include files
+#include "FWCore/Framework/interface/ModuleFactory.h"
+#include "FWCore/Framework/interface/ESProducer.h"
+
+#include "FWCore/Framework/interface/ESHandle.h"
+#include "TrackingTools/Records/interface/TfGraphRecord.h"
+#include "DataFormats/TrackTfGraph/interface/TfGraphDefWrapper.h"
+
+// class declaration
+
+class TfGraphDefProducer : public edm::ESProducer {
+public:
+  TfGraphDefProducer(const edm::ParameterSet&);
+
+  using ReturnType = std::unique_ptr<TfGraphDefWrapper>;
+
+  ReturnType produce(const TfGraphRecord&);
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+    TfGraphDefWrapper wrapper_;
+
+  // ----------member data ---------------------------
+};
+
+TfGraphDefProducer::TfGraphDefProducer(const edm::ParameterSet& iConfig):
+wrapper_(TfGraphDefWrapper(tensorflow::loadGraphDef(iConfig.getParameter<edm::FileInPath>("FileName").fullPath())))
+{
+  auto componentName = iConfig.getParameter<std::string>("ComponentName");
+  setWhatProduced(this, componentName);
+}
+
+// ------------ method called to produce the data  ------------
+std::unique_ptr<TfGraphDefWrapper> TfGraphDefProducer::produce(const TfGraphRecord& iRecord) {
+  return std::unique_ptr<TfGraphDefWrapper>(&wrapper_);
+}
+
+void TfGraphDefProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<std::string>("ComponentName", "tfGraphDef");
+  desc.add<edm::FileInPath>("FileName", edm::FileInPath());
+  descriptions.add("tfGraphDefProducer", desc);
+}
+
+//define this as a plug-in
+#include "FWCore/PluginManager/interface/ModuleDef.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+DEFINE_FWK_EVENTSETUP_MODULE(TfGraphDefProducer);

DataFormats/TrackTfGraph/src/ES_TfGraphDefWrapper.cc

@@ -0,0 +1,4 @@
+#include "DataFormats/TrackTfGraph/interface/TfGraphDefWrapper.h"
+#include "FWCore/Utilities/interface/typelookup.h"
+
+TYPELOOKUP_DATA_REG(TfGraphDefWrapper);

DataFormats/TrackTfGraph/src/TfGraphDefWrapper.cc

@@ -0,0 +1,7 @@
+#include "DataFormats/TrackTfGraph/interface/TfGraphDefWrapper.h"
+
+TfGraphDefWrapper::TfGraphDefWrapper(tensorflow::GraphDef* graph) {graphDef_ = graph;}
+
+tensorflow::GraphDef* TfGraphDefWrapper::GetGraphDef() const {
+    return graphDef_;
+}

RecoTracker/FinalTrackSelectors/BuildFile.xml

@@ -7,6 +7,8 @@
 <use name="CondFormats/EgammaObjects"/>
 <use name="roottmva"/>
 <use name="lwtnn"/>
+<use name="TrackingTools/Records"/>
+<use name="PhysicsTools/TensorFlow"/>
 <export>
   <lib name="1"/>
 </export>

RecoTracker/FinalTrackSelectors/plugins/BuildFile.xml

@@ -34,6 +34,7 @@
 <use name="clhep"/>
 <use name="roottmva"/>
 <use name="lwtnn"/>
+<use name="DataFormats/TrackTfGraph"/>
 <library file="*.cc" name="RecoTrackerFinalTrackSelectorsPlugins">
   <flags EDM_PLUGIN="1"/>
 </library>

RecoTracker/FinalTrackSelectors/plugins/TrackTfClassifier.cc

@@ -0,0 +1,122 @@
+#include "RecoTracker/FinalTrackSelectors/interface/TrackMVAClassifier.h"
+
+#include "FWCore/Framework/interface/EventSetup.h"
+#include "FWCore/Framework/interface/ESHandle.h"
+
+#include "DataFormats/TrackReco/interface/Track.h"
+#include "DataFormats/VertexReco/interface/Vertex.h"
+
+#include "getBestVertex.h"
+
+#include "TrackingTools/Records/interface/TfGraphRecord.h"
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include "DataFormats/TrackTfGraph/interface/TfGraphDefWrapper.h"
+
+struct TfDnnCache {
+    TfDnnCache() : session_(nullptr) {}
+    tensorflow::Session* session_;
+};
+
+namespace {
+    struct tfDnn {
+        tfDnn(const edm::ParameterSet& cfg):
+            tfDnnLabel_(cfg.getParameter<std::string>("tfDnnLabel"))
+
+        {}
+        TfDnnCache* cache_ = new TfDnnCache();
+
+        ~tfDnn() {
+            if(cache_->session_) {
+                tensorflow::closeSession(cache_->session_);
+                delete cache_;
+            }
+        }
+
+        static const char *name() { return "TrackTfClassifier"; }
+
+        static void fillDescriptions(edm::ParameterSetDescription& desc) {
+            desc.add<std::string>("tfDnnLabel", "trackSelectionTf");
+        }
+
+        void beginStream() {}
+
+        void initEvent(const edm::EventSetup& es) {
+            if (!cache_->session_) {
+                edm::ESHandle<TfGraphDefWrapper> tfDnnHandle;
+                es.get<TfGraphRecord>().get(tfDnnLabel_, tfDnnHandle);
+                tensorflow::GraphDef *graphDef_ = tfDnnHandle.product()->GetGraphDef();
+                cache_->session_ = tensorflow::createSession(graphDef_, 1); //The integer controls how many threads are used for running inference
+            }
+            session_ = cache_->session_;
+        }
+
+        float operator()(reco::Track const & trk,
+                        reco::BeamSpot const & beamSpot,
+                        reco::VertexCollection const & vertices) const {
+
+            Point bestVertex = getBestVertex(trk, vertices);
+
+            tensorflow::Tensor input1(tensorflow::DT_FLOAT, {1, 29});
+            tensorflow::Tensor input2(tensorflow::DT_FLOAT, {1, 1});
+
+            input1.matrix<float>()(0, 0) = trk.pt();
+            input1.matrix<float>()(0, 1) = trk.innerMomentum().x();
+            input1.matrix<float>()(0, 2) = trk.innerMomentum().y();
+            input1.matrix<float>()(0, 3) = trk.innerMomentum().z();
+            input1.matrix<float>()(0, 4) = trk.innerMomentum().rho();
+            input1.matrix<float>()(0, 5) = trk.outerMomentum().x();
+            input1.matrix<float>()(0, 6) = trk.outerMomentum().y();
+            input1.matrix<float>()(0, 7) = trk.outerMomentum().z();
+            input1.matrix<float>()(0, 8) = trk.outerMomentum().rho();
+            input1.matrix<float>()(0, 9) = trk.ptError();
+            input1.matrix<float>()(0, 10) = trk.dxy(bestVertex);
+            input1.matrix<float>()(0, 11) = trk.dz(bestVertex);
+            input1.matrix<float>()(0, 12) = trk.dxy(beamSpot.position());
+            input1.matrix<float>()(0, 13) = trk.dz(beamSpot.position());
+            input1.matrix<float>()(0, 14) = trk.dxyError();
+            input1.matrix<float>()(0, 15) = trk.dzError();
+            input1.matrix<float>()(0, 16) = trk.normalizedChi2();
+            input1.matrix<float>()(0, 17) = trk.eta();
+            input1.matrix<float>()(0, 18) = trk.phi();
+            input1.matrix<float>()(0, 19) = trk.etaError();
+            input1.matrix<float>()(0, 20) = trk.phiError();
+            input1.matrix<float>()(0, 21) = trk.hitPattern().numberOfValidPixelHits();
+            input1.matrix<float>()(0, 22) = trk.hitPattern().numberOfValidStripHits();
+            input1.matrix<float>()(0, 23) = trk.ndof();
+            input1.matrix<float>()(0, 24) = trk.hitPattern().numberOfLostTrackerHits(reco::HitPattern::MISSING_INNER_HITS);
+            input1.matrix<float>()(0, 25) = trk.hitPattern().numberOfLostTrackerHits(reco::HitPattern::MISSING_OUTER_HITS);
+            input1.matrix<float>()(0, 26) = trk.hitPattern().trackerLayersTotallyOffOrBad(reco::HitPattern::MISSING_INNER_HITS);
+            input1.matrix<float>()(0, 27) = trk.hitPattern().trackerLayersTotallyOffOrBad(reco::HitPattern::MISSING_OUTER_HITS);
+            input1.matrix<float>()(0, 28) = trk.hitPattern().trackerLayersWithoutMeasurement(reco::HitPattern::TRACK_HITS);
+
+            //Original algo as its own input, it will enter the graph so that it gets one-hot encoded, as is the preferred
+            //format for categorical inputs, where the labels do not have any metric amongst them
+            input2.matrix<float>()(0, 0) = trk.originalAlgo();
+
+            //The names for the input tensors get locked when freezing the trained tensorflow model. The NamedTensors must
+            //match those names
+            tensorflow::NamedTensorList inputs;
+            inputs.resize(2);
+            inputs[0] = tensorflow::NamedTensor("x", input1);
+            inputs[1] = tensorflow::NamedTensor("y", input2);
+            std::vector<tensorflow::Tensor> outputs;
+
+            //evaluate the input
+            tensorflow::run(session_, inputs, { "Identity" }, &outputs);
+
+            //scale output to be [-1, 1] due to convention
+            float output = 2.0*outputs[0].matrix<float>()(0, 0)-1.0;
+            return output;
+        }
+
+        std::string tfDnnLabel_;
+        tensorflow::Session *session_;
+    };
+
+    using TrackTfClassifier = TrackMVAClassifier<tfDnn>;
+}
+#include "FWCore/PluginManager/interface/ModuleDef.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+DEFINE_FWK_MODULE(TrackTfClassifier);
+

RecoTracker/FinalTrackSelectors/python/trackSelectionTf_cfi.py

@@ -0,0 +1,5 @@
+from DataFormats.TrackTfGraph.tfGraphDefProducer_cfi import tfGraphDefProducer as _tfGraphDefProducer
+trackSelectionTf = _tfGraphDefProducer.clone(
+    ComponentName = "trackSelectionTf",
+    FileName = "RecoTracker/FinalTrackSelectors/data/frozen_graph.pb"
+)

RecoTracker/FinalTrackSelectors/src/ES_Tf.cc

@@ -0,0 +1,3 @@
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include "FWCore/Utilities/interface/typelookup.h"
+TYPELOOKUP_DATA_REG(tensorflow::Session);

RecoTracker/IterativeTracking/python/DetachedQuadStep_cff.py

@@ -4,6 +4,7 @@
 
 #for dnn classifier
 from Configuration.ProcessModifiers.trackdnn_cff import trackdnn
+from dnnQualityCuts import qualityCutDictionary
 
 ###############################################
 # Low pT and detached tracks from pixel quadruplets
@@ -212,11 +213,11 @@
     qualityCuts = [-0.5,0.0,0.5]
 )
 
-from RecoTracker.FinalTrackSelectors.TrackLwtnnClassifier_cfi import *
-from RecoTracker.FinalTrackSelectors.trackSelectionLwtnn_cfi import *
-trackdnn.toReplaceWith(detachedQuadStep, TrackLwtnnClassifier.clone(
-    src         = 'detachedQuadStepTracks',
-    qualityCuts = [-0.6, 0.05, 0.7]
+from RecoTracker.FinalTrackSelectors.TrackTfClassifier_cfi import *
+from RecoTracker.FinalTrackSelectors.trackSelectionTf_cfi import *
+trackdnn.toReplaceWith(detachedQuadStep, TrackTfClassifier.clone(
+    src = 'detachedQuadStepTracks',
+    qualityCuts = qualityCutDictionary['DetachedQuadStep']
 ))
 
 highBetaStar_2018.toModify(detachedQuadStep,qualityCuts = [-0.7,0.0,0.5])

RecoTracker/IterativeTracking/python/DetachedTripletStep_cff.py

@@ -8,6 +8,7 @@
 
 #for dnn classifier
 from Configuration.ProcessModifiers.trackdnn_cff import trackdnn
+from dnnQualityCuts import qualityCutDictionary
 
 ###############################################
 # Low pT and detached tracks from pixel triplets
@@ -260,11 +261,11 @@
     qualityCuts = [-0.2,0.3,0.8]
 ))
 
-from RecoTracker.FinalTrackSelectors.TrackLwtnnClassifier_cfi import *
-from RecoTracker.FinalTrackSelectors.trackSelectionLwtnn_cfi import *
-trackdnn.toReplaceWith(detachedTripletStep, TrackLwtnnClassifier.clone(
+from RecoTracker.FinalTrackSelectors.TrackTfClassifier_cfi import *
+from RecoTracker.FinalTrackSelectors.trackSelectionTf_cfi import *
+trackdnn.toReplaceWith(detachedTripletStep, TrackTfClassifier.clone(
      src = 'detachedTripletStepTracks',
-     qualityCuts = [0.0, 0.4, 0.8],
+     qualityCuts = qualityCutDictionary['DetachedTripletStep'],
 ))
 (trackdnn & fastSim).toModify(detachedTripletStep,vertices = 'firstStepPrimaryVerticesBeforeMixing')
 

RecoTracker/IterativeTracking/python/HighPtTripletStep_cff.py

@@ -4,6 +4,7 @@
 
 #for dnn classifier
 from Configuration.ProcessModifiers.trackdnn_cff import trackdnn
+from dnnQualityCuts import qualityCutDictionary
 
 ### high-pT triplets ###
 
@@ -254,11 +255,11 @@
      qualityCuts = [0.2,0.3,0.4]
 )
 
-from RecoTracker.FinalTrackSelectors.TrackLwtnnClassifier_cfi import *
-from RecoTracker.FinalTrackSelectors.trackSelectionLwtnn_cfi import *
-trackdnn.toReplaceWith(highPtTripletStep, TrackLwtnnClassifier.clone(
-    src         = 'highPtTripletStepTracks',
-    qualityCuts = [0.75, 0.775, 0.8],
+from RecoTracker.FinalTrackSelectors.TrackTfClassifier_cfi import *
+from RecoTracker.FinalTrackSelectors.trackSelectionTf_cfi import *
+trackdnn.toReplaceWith(highPtTripletStep, TrackTfClassifier.clone(
+    src = 'highPtTripletStepTracks',
+    qualityCuts = qualityCutDictionary['HighPtTripletStep'],
 ))
 
 highBetaStar_2018.toModify(highPtTripletStep,qualityCuts = [-0.2,0.3,0.4])

RecoTracker/IterativeTracking/python/InitialStep_cff.py

@@ -4,6 +4,7 @@
 
 #for dnn classifier
 from Configuration.ProcessModifiers.trackdnn_cff import trackdnn
+from dnnQualityCuts import qualityCutDictionary
 
 ### STEP 0 ###
 
@@ -318,13 +319,14 @@
      qualityCuts = [-0.95,-0.85,-0.75]
 ))
 
-from RecoTracker.FinalTrackSelectors.TrackLwtnnClassifier_cfi import *
-from RecoTracker.FinalTrackSelectors.trackSelectionLwtnn_cfi import *
-trackdnn.toReplaceWith(initialStep, TrackLwtnnClassifier.clone(
-        src         = 'initialStepTracks',
-        qualityCuts = [0.0, 0.3, 0.6]
+from RecoTracker.FinalTrackSelectors.TrackTfClassifier_cfi import *
+from RecoTracker.FinalTrackSelectors.trackSelectionTf_cfi import *
+trackdnn.toReplaceWith(initialStep, TrackTfClassifier.clone(
+        src = 'initialStepTracks',
+        qualityCuts = qualityCutDictionary["InitialStep"]
 ))
-(trackdnn & fastSim).toModify(initialStep,vertices = 'firstStepPrimaryVerticesBeforeMixing')
+
+(trackdnn & fastSim).toModify(initialStep,vertices = "firstStepPrimaryVerticesBeforeMixing")
 
 pp_on_AA_2018.toModify(initialStep, 
         mva         = dict(GBRForestLabel = 'HIMVASelectorInitialStep_Phase1'),

RecoTracker/IterativeTracking/python/JetCoreRegionalStep_cff.py

@@ -2,6 +2,7 @@
 
 #for dnn classifier
 from Configuration.ProcessModifiers.trackdnn_cff import trackdnn
+from dnnQualityCuts import qualityCutDictionary
 
 # This step runs over all clusters
 
@@ -198,11 +199,11 @@
      qualityCuts = [-0.2,0.0,0.4]
 ))
 
-from RecoTracker.FinalTrackSelectors.TrackLwtnnClassifier_cfi import *
-from RecoTracker.FinalTrackSelectors.trackSelectionLwtnn_cfi import *
-trackdnn.toReplaceWith(jetCoreRegionalStep, TrackLwtnnClassifier.clone(
-     src         = 'jetCoreRegionalStepTracks',
-     qualityCuts = [0.6, 0.7, 0.8],
+from RecoTracker.FinalTrackSelectors.TrackTfClassifier_cfi import *
+from RecoTracker.FinalTrackSelectors.trackSelectionTf_cfi import *
+trackdnn.toReplaceWith(jetCoreRegionalStep, TrackTfClassifier.clone(
+     src = 'jetCoreRegionalStepTracks',
+     qualityCuts = qualityCutDictionary["JetCoreRegionalStep"],
 ))
 
 fastSim.toModify(jetCoreRegionalStep,vertices = 'firstStepPrimaryVerticesBeforeMixing')