Track Purity DNN for Phase-2 HLT

Configuration/ProcessModifiers/python/trackTorchClassifier_cff.py

@@ -0,0 +1,4 @@
+import FWCore.ParameterSet.Config as cms
+
+# This modifier sets the use of a deep neural network for high purity track selection  
+trackTorchClassifier = cms.Modifier()

HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py

@@ -47,3 +47,9 @@
 hltPhase2LegacyTracking.toModify(hltInitialStepTrackCutClassifier,
     mva = dict(passThroughForAll=False, passThroughForDisplaced=False)
 )
+
+
+from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier
+trackTorchClassifier.toModify(hltInitialStepTrackCutClassifier,
+    src = "hltInitialStepTrackTorchClassifierOutput"
+)

HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackFeatureExtractor_cfi.py

@@ -0,0 +1,9 @@
+import FWCore.ParameterSet.Config as cms
+
+hltInitialStepTrackFeatureExtractor = cms.EDProducer("TrackFeatureExtractor@alpaka",
+    src = cms.InputTag("hltInitialStepTracks"),
+    beamSpot = cms.InputTag("hltOnlineBeamSpot"),
+    alpaka = cms.untracked.PSet(
+        backend = cms.untracked.string('')
+    )
+)

HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py

@@ -8,3 +8,6 @@
     originalQualVals = cms.InputTag("hltInitialStepTrackCutClassifier","QualityMasks"),
     originalSource = cms.InputTag("hltInitialStepTracks")
 )
+
+from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier
+trackTorchClassifier.toModify(hltInitialStepTrackSelectionHighPurity, originalSource = "hltInitialStepTrackTorchClassifierOutput")

HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifierOutput_cfi.py

@@ -0,0 +1,10 @@
+import FWCore.ParameterSet.Config as cms
+
+hltInitialStepTrackTorchClassifierOutput = cms.EDProducer("TrackTorchClassifierFromSoA",
+    src = cms.InputTag("hltInitialStepTracks"),
+    scores = cms.InputTag("hltInitialStepTrackTorchClassifier"),
+    features = cms.InputTag("hltInitialStepTrackFeatureExtractor"),
+    minScore = cms.double(0.377),
+    dxyThreshold = cms.double(0.5),
+    highDxyMinScore = cms.double(0.267)
+)

HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifier_cfi.py

@@ -0,0 +1,9 @@
+import FWCore.ParameterSet.Config as cms
+
+hltInitialStepTrackTorchClassifier = cms.EDProducer("TrackTorchClassifierAlpaka@alpaka",
+    modelPath = cms.FileInPath('RecoTracker/FinalTrackSelectors/data/TrackTorchClassifier/model.pt'),
+    features = cms.InputTag("hltInitialStepTrackFeatureExtractor"),
+    alpaka = cms.untracked.PSet(
+        backend = cms.untracked.string('')
+    )
+)

HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepHPSelectionSequence_cfi.py

@@ -0,0 +1,9 @@
+import FWCore.ParameterSet.Config as cms
+
+from ..modules.hltInitialStepTrackCutClassifier_cfi import *
+from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import *
+
+HLTInitialStepHPSelectionSequence = cms.Sequence(
+    hltInitialStepTrackCutClassifier
+    +hltInitialStepTrackSelectionHighPurity
+)

HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py

@@ -7,14 +7,14 @@
 from ..modules.hltInitialStepSeeds_cfi import *
 from ..modules.hltInitialStepTrackCandidates_cfi import *
 from ..modules.hltInitialStepTrackCandidatesMkFit_cfi import *
-from ..modules.hltInitialStepTrackCutClassifier_cfi import *
-from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import *
 from ..modules.hltInitialStepTracks_cfi import *
 from ..modules.hltInitialStepTrajectorySeedsLST_cfi import *
 from ..modules.hltInitialStepTrajectorySeedsLSTTracks_cfi import *
 from ..modules.hltLST_cfi import *
 from ..modules.hltSiPhase2RecHits_cfi import *
 from ..sequences.HLTMkFitInputSequence_cfi import *
+from ..sequences.HLTInitialStepHPSelectionSequence_cfi import *
+from ..sequences.HLTInitialStepTorchClassifierHPSelectionSequence_cfi import *
 
 HLTInitialStepSequence = cms.Sequence(
      hltInitialStepSeeds
@@ -28,8 +28,7 @@
     +hltInitialStepTrackCandidatesMkFit
     +hltInitialStepTrackCandidates
     +hltInitialStepTracks
-    +hltInitialStepTrackCutClassifier
-    +hltInitialStepTrackSelectionHighPurity
+    +HLTInitialStepHPSelectionSequence
 )
 
 
@@ -97,8 +96,7 @@
     +hltLST
     +hltInitialStepTrackCandidates
     +hltInitialStepTracks
-    +hltInitialStepTrackCutClassifier
-    +hltInitialStepTrackSelectionHighPurity
+    +HLTInitialStepHPSelectionSequence
 )
 
 from Configuration.ProcessModifiers.trackingLST_cff import trackingLST
@@ -133,9 +131,18 @@
     +hltInitialStepTrackCandidatesMkFit
     +hltInitialStepTrackCandidatesMkFitFit
     +hltInitialStepTracks
-    +hltInitialStepTrackCutClassifier
-    +hltInitialStepTrackSelectionHighPurity
+    +HLTInitialStepHPSelectionSequence
 )
 
 from Configuration.ProcessModifiers.trackingMkFitFit_cff import trackingMkFitFit
 trackingMkFitFit.toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceMkFitFit)
+
+_HLTInitialStepSequenceTrackTorchClassifier = HLTInitialStepSequence.copyAndExclude([HLTInitialStepHPSelectionSequence]) 
+_HLTInitialStepSequenceTrackTorchClassifier += HLTInitialStepTorchClassifierHPSelectionSequence 
+
+_HLTInitialStepSequenceTrackTorchClassifierMkFitFit = _HLTInitialStepSequenceMkFitFit.copyAndExclude([HLTInitialStepHPSelectionSequence])
+_HLTInitialStepSequenceTrackTorchClassifierMkFitFit += HLTInitialStepTorchClassifierHPSelectionSequence 
+
+from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier
+trackTorchClassifier.toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackTorchClassifier)
+(trackTorchClassifier & trackingMkFitFit).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackTorchClassifierMkFitFit)

HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepTorchClassifierHPSelectionSequence_cfi.py

@@ -0,0 +1,15 @@
+import FWCore.ParameterSet.Config as cms
+
+from ..modules.hltInitialStepTrackCutClassifier_cfi import *
+from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import *
+from ..modules.hltInitialStepTrackFeatureExtractor_cfi import *
+from ..modules.hltInitialStepTrackTorchClassifier_cfi import *
+from ..modules.hltInitialStepTrackTorchClassifierOutput_cfi import *
+
+HLTInitialStepTorchClassifierHPSelectionSequence = cms.Sequence(
+    hltInitialStepTrackFeatureExtractor
+    +hltInitialStepTrackTorchClassifier
+    +hltInitialStepTrackTorchClassifierOutput 
+    +hltInitialStepTrackCutClassifier
+    +hltInitialStepTrackSelectionHighPurity
+)

RecoTracker/FinalTrackSelectors/BuildFile.xml

@@ -1,13 +1,20 @@
 <use name="ofast-flag"/>
+<use name="alpaka"/>
+<use name="DataFormats/Common"/>
+<use name="DataFormats/Portable"/>
+<use name="DataFormats/SoATemplate"/>
 <use name="DataFormats/TrackReco"/>
 <use name="DataFormats/VertexReco"/>
 <use name="FWCore/Framework"/>
 <use name="FWCore/MessageLogger"/>
 <use name="FWCore/ParameterSet"/>
 <use name="CondFormats/GBRForest"/>
 <use name="TrackingTools/PatternTools"/>
+<use name="HeterogeneousCore/AlpakaInterface"/>
 <use name="roottmva"/>
 <use name="lwtnn"/>
+<use name="clhep"/>
+<flags ALPAKA_BACKENDS="1"/>
 <export>
   <lib name="1"/>
 </export>

RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h

@@ -0,0 +1,30 @@
+#ifndef RecoTracker_FinalTrackSelectors_TrackTorchClassifierFeaturesSoA_h
+#define RecoTracker_FinalTrackSelectors_TrackTorchClassifierFeaturesSoA_h
+
+#include "DataFormats/SoATemplate/interface/SoALayout.h"
+
+GENERATE_SOA_LAYOUT(TrackTorchClassifierFeaturesSoALayout,
+                    SOA_COLUMN(float, dxyBeamSpot),
+                    SOA_COLUMN(float, dzBeamSpot),
+                    SOA_COLUMN(float, dxyError),
+                    SOA_COLUMN(float, dzError),
+                    SOA_COLUMN(float, normalizedChi2),
+                    SOA_COLUMN(float, eta),
+                    SOA_COLUMN(float, phi),
+                    SOA_COLUMN(float, etaError),
+                    SOA_COLUMN(float, phiError),
+                    SOA_COLUMN(float, ndof),
+                    SOA_COLUMN(float, lostInnerHits),
+                    SOA_COLUMN(float, lostOuterHits),
+                    SOA_COLUMN(float, layersWithoutMeas),
+                    SOA_COLUMN(float, validPixelHits),
+                    SOA_COLUMN(float, validStripHits))
+
+using TrackTorchClassifierFeaturesSoA = TrackTorchClassifierFeaturesSoALayout<>;
+
+// Define the SoA layout for track scores (output)
+GENERATE_SOA_LAYOUT(TrackTorchClassifierScoresSoALayout, SOA_COLUMN(float, score))
+
+using TrackTorchClassifierScoresSoA = TrackTorchClassifierScoresSoALayout<>;
+
+#endif  // RecoTracker_FinalTrackSelectors_TrackTorchClassifierFeaturesSoA_h

RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h

@@ -0,0 +1,12 @@
+#ifndef RecoTracker_FinalTrackSelectors_alpaka_TrackFeaturesDeviceCollection_h
+#define RecoTracker_FinalTrackSelectors_alpaka_TrackFeaturesDeviceCollection_h
+
+#include "DataFormats/Portable/interface/alpaka/PortableCollection.h"
+#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h"
+#include "HeterogeneousCore/AlpakaInterface/interface/config.h"
+
+namespace ALPAKA_ACCELERATOR_NAMESPACE {
+  using TrackFeaturesDeviceCollection = PortableCollection<TrackTorchClassifierFeaturesSoA>;
+}  // namespace ALPAKA_ACCELERATOR_NAMESPACE
+
+#endif  // RecoTracker_FinalTrackSelectors_alpaka_TrackFeaturesDeviceCollection_h

RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h

@@ -0,0 +1,12 @@
+#ifndef RecoTracker_FinalTrackSelectors_alpaka_TrackScoresDeviceCollection_h
+#define RecoTracker_FinalTrackSelectors_alpaka_TrackScoresDeviceCollection_h
+
+#include "DataFormats/Portable/interface/alpaka/PortableCollection.h"
+#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h"
+#include "HeterogeneousCore/AlpakaInterface/interface/config.h"
+
+namespace ALPAKA_ACCELERATOR_NAMESPACE {
+  using TrackScoresDeviceCollection = PortableCollection<TrackTorchClassifierScoresSoA>;
+}  // namespace ALPAKA_ACCELERATOR_NAMESPACE
+
+#endif  // RecoTracker_FinalTrackSelectors_alpaka_TrackScoresDeviceCollection_h

RecoTracker/FinalTrackSelectors/plugins/BuildFile.xml

@@ -24,6 +24,7 @@
 <use name="MagneticField/Engine"/>
 <use name="MagneticField/Records"/>
 <use name="PhysicsTools/TensorFlow"/>
+<use name="PhysicsTools/PyTorch"/>
 <use name="RecoLocalTracker/SiStripClusterizer"/>
 <use name="RecoTracker/FinalTrackSelectors"/>
 <use name="RecoTracker/Record"/>
@@ -38,3 +39,16 @@
 <library file="*.cc" name="RecoTrackerFinalTrackSelectorsPlugins">
   <flags EDM_PLUGIN="1"/>
 </library>
+
+<library file="alpaka/*.cc" name="RecoTrackerFinalTrackSelectorsAlpakaPlugins">
+  <use name="alpaka"/>
+  <use name="DataFormats/BeamSpot"/>
+  <use name="DataFormats/Portable"/>
+  <use name="DataFormats/SoATemplate"/>
+  <use name="HeterogeneousCore/AlpakaCore"/>
+  <use name="HeterogeneousCore/AlpakaInterface"/>
+  <use name="PhysicsTools/PyTorchAlpaka"/>
+  <use name="pytorch-cuda" for="alpaka/cuda"/>
+  <flags ALPAKA_BACKENDS="1"/>
+  <flags EDM_PLUGIN="1"/>
+</library>

RecoTracker/FinalTrackSelectors/plugins/TrackTorchClassifierFromSoA.cc

@@ -0,0 +1,88 @@
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+
+#include "DataFormats/TrackReco/interface/Track.h"
+#include "DataFormats/TrackReco/interface/TrackFwd.h"
+
+#include "DataFormats/Portable/interface/PortableHostCollection.h"
+#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h"
+
+// This module consumes the HOST copy of the Alpaka device scores
+// The framework automatically creates host copies of device PortableCollections
+class TrackTorchClassifierFromSoA : public edm::stream::EDProducer<> {
+public:
+  explicit TrackTorchClassifierFromSoA(const edm::ParameterSet& iConfig);
+  ~TrackTorchClassifierFromSoA() override = default;
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+  void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
+
+  const edm::EDGetTokenT<reco::TrackCollection> tracks_token_;
+  const edm::EDGetTokenT<PortableHostCollection<TrackTorchClassifierScoresSoA>> scores_token_;
+  const edm::EDGetTokenT<PortableHostCollection<TrackTorchClassifierFeaturesSoA>> features_token_;
+  const float min_score_;
+  const float dxy_threshold_;
+  const float high_dxy_min_score_;
+
+  const edm::EDPutTokenT<reco::TrackCollection> filtered_tracks_token_;
+  const edm::EDPutTokenT<std::vector<float>> scores_output_token_;
+};
+
+TrackTorchClassifierFromSoA::TrackTorchClassifierFromSoA(const edm::ParameterSet& iConfig)
+    : tracks_token_(consumes(iConfig.getParameter<edm::InputTag>("src"))),
+      scores_token_(consumes(iConfig.getParameter<edm::InputTag>("scores"))),
+      features_token_(consumes(iConfig.getParameter<edm::InputTag>("features"))),
+      min_score_(iConfig.getParameter<double>("minScore")),
+      dxy_threshold_(iConfig.getParameter<double>("dxyThreshold")),
+      high_dxy_min_score_(iConfig.getParameter<double>("highDxyMinScore")),
+      filtered_tracks_token_(produces<reco::TrackCollection>()),
+      scores_output_token_(produces<std::vector<float>>("MVAScores")) {}
+
+void TrackTorchClassifierFromSoA::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("src", edm::InputTag("hltInitialStepTracks"));
+  desc.add<edm::InputTag>("scores", edm::InputTag("hltInitialStepTrackTorchClassifier"));
+  desc.add<edm::InputTag>("features", edm::InputTag("hltInitialStepTrackTorchClassifier"));
+  desc.add<double>("minScore", 0.5)->setComment("Minimum DNN score to keep track (working point)");
+  desc.add<double>("dxyThreshold", 0.5)->setComment("Tracks with |dxy| > this value bypass the score cut");
+  desc.add<double>("highDxyMinScore", 0.5)->setComment("Minimum DNN score to keep high dxy track (working point)");
+  descriptions.addWithDefaultLabel(desc);
+}
+
+void TrackTorchClassifierFromSoA::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  const auto& tracks = iEvent.get(tracks_token_);
+  const auto& scores_host = iEvent.get(scores_token_);
+  const auto& features_host = iEvent.get(features_token_);
+
+  const auto nTracks = tracks.size();
+
+  // Create filtered track collection
+  auto filtered_tracks = std::make_unique<reco::TrackCollection>();
+  auto all_scores = std::make_unique<std::vector<float>>();
+  all_scores->reserve(nTracks);
+
+  // Access scores and features from the host collection
+  auto scores_view = scores_host.const_view();
+  auto features_view = features_host.const_view();
+
+  for (size_t i = 0; i < nTracks; ++i) {
+    float score = scores_view[i].score();
+    float dxy = features_view[i].dxyBeamSpot();
+    all_scores->push_back(score);
+
+    if (score >= min_score_ || ((std::abs(dxy) > dxy_threshold_) && (score >= high_dxy_min_score_))) {
+      filtered_tracks->push_back(tracks[i]);
+    }
+  }
+
+  iEvent.put(filtered_tracks_token_, std::move(filtered_tracks));
+  iEvent.put(scores_output_token_, std::move(all_scores));
+}
+
+DEFINE_FWK_MODULE(TrackTorchClassifierFromSoA);