Track Purity DNN for Phase-2 HLT

Configuration/ProcessModifiers/python/trackCutClassifier_cff.py

@@ -0,0 +1,4 @@
+import FWCore.ParameterSet.Config as cms
+
+# This modifier sets the use of cut based high purity track selection  
+trackCutClassifier = cms.Modifier()

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

@@ -32,7 +32,7 @@
         minLayersForDisplaced = cms.int32(4)
     ),
     qualityCuts = cms.vdouble(-0.7, 0.1, 0.7),
-    src = cms.InputTag("hltInitialStepTracks"),
+    src = cms.InputTag("hltInitialStepTrackTorchClassifierOutput"),
     vertices = cms.InputTag("hltPhase2PixelVertices")
 )
 
@@ -47,3 +47,9 @@
 hltPhase2LegacyTracking.toModify(hltInitialStepTrackCutClassifier,
     mva = dict(passThroughForAll=False, passThroughForDisplaced=False)
 )
+
+
+from Configuration.ProcessModifiers.trackCutClassifier_cff import trackCutClassifier
+trackCutClassifier.toModify(hltInitialStepTrackCutClassifier,
+    src = "hltInitialStepTracks"
+)

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

@@ -0,0 +1,9 @@
+import FWCore.ParameterSet.Config as cms
+
+hltInitialStepTrackFeatureExtractor = cms.EDProducer("alpaka_serial_sync::TrackFeatureExtractor",
+    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

@@ -6,8 +6,11 @@
     minQuality = cms.string('highPurity'),
     originalMVAVals = cms.InputTag("hltInitialStepTrackCutClassifier","MVAValues"),
     originalQualVals = cms.InputTag("hltInitialStepTrackCutClassifier","QualityMasks"),
-    originalSource = cms.InputTag("hltInitialStepTracks")
+    originalSource = cms.InputTag("hltInitialStepTrackTorchClassifierOutput")
 )
 
 from Configuration.ProcessModifiers.mtd_at_hlt_cff import mtd_at_hlt
 mtd_at_hlt.toModify(hltInitialStepTrackSelectionHighPurity, copyTrajectories = True)
+
+from Configuration.ProcessModifiers.trackCutClassifier_cff import trackCutClassifier
+trackCutClassifier.toModify(hltInitialStepTrackSelectionHighPurity, originalSource = "hltInitialStepTracks")

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

@@ -0,0 +1,14 @@
+import FWCore.ParameterSet.Config as cms
+
+hltInitialStepTrackTorchClassifierOutput = cms.EDProducer("TrackTorchClassifierFromSoA",
+    src = cms.InputTag("hltInitialStepTracks"),
+    scores = cms.InputTag("hltInitialStepTrackTorchClassifier"),
+    features = cms.InputTag("hltInitialStepTrackFeatureExtractor"),
+    copyTrajectories = cms.bool(False),
+    minScore = cms.double(0.377),
+    dxyThreshold = cms.double(0.5),
+    highDxyMinScore = cms.double(0.267)
+)
+
+from Configuration.ProcessModifiers.mtd_at_hlt_cff import mtd_at_hlt
+mtd_at_hlt.toModify(hltInitialStepTrackTorchClassifierOutput, copyTrajectories = True)

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

@@ -0,0 +1,9 @@
+import FWCore.ParameterSet.Config as cms
+
+hltInitialStepTrackTorchClassifier = cms.EDProducer("alpaka_serial_sync::TrackTorchClassifierAlpaka",
+    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/HLTInitialStepCutClassifierHPSelectionSequence_cfi.py

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

HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepHPSelectionSequence_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 *
+
+HLTInitialStepHPSelectionSequence = cms.Sequence(
+    hltInitialStepTrackFeatureExtractor
+    +hltInitialStepTrackTorchClassifier
+    +hltInitialStepTrackTorchClassifierOutput 
+    +hltInitialStepTrackCutClassifier
+    +hltInitialStepTrackSelectionHighPurity
+)

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

@@ -6,13 +6,13 @@
 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 ..sequences.HLTMkFitInputSequence_cfi import *
+from .HLTInitialStepHPSelectionSequence_cfi import *
+from .HLTInitialStepCutClassifierHPSelectionSequence_cfi import *
 
 HLTInitialStepSequence = cms.Sequence(
      hltInitialStepSeeds
@@ -24,8 +24,7 @@
     +hltInitialStepTrackCandidatesMkFit
     +hltInitialStepTrackCandidates
     +hltInitialStepTracks
-    +hltInitialStepTrackCutClassifier
-    +hltInitialStepTrackSelectionHighPurity
+    +HLTInitialStepHPSelectionSequence
 )
 
 
@@ -84,8 +83,7 @@
     +hltLST
     +hltInitialStepTrackCandidates
     +hltInitialStepTracks
-    +hltInitialStepTrackCutClassifier
-    +hltInitialStepTrackSelectionHighPurity
+    +HLTInitialStepHPSelectionSequence
 )
 
 from Configuration.ProcessModifiers.trackingLST_cff import trackingLST
@@ -116,9 +114,18 @@
     +hltInitialStepTrackCandidatesMkFit
     +hltInitialStepTrackCandidatesMkFitFit
     +hltInitialStepTracks
-    +hltInitialStepTrackCutClassifier
-    +hltInitialStepTrackSelectionHighPurity
+    +HLTInitialStepHPSelectionSequence
 )
 
 from Configuration.ProcessModifiers.trackingMkFitFit_cff import trackingMkFitFit
 trackingMkFitFit.toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceMkFitFit)
+
+_HLTInitialStepSequenceTrackCutClassifier = HLTInitialStepSequence.copyAndExclude([HLTInitialStepHPSelectionSequence]) 
+_HLTInitialStepSequenceTrackCutClassifier += HLTInitialStepCutClassifierHPSelectionSequence 
+
+_HLTInitialStepSequenceTrackCutClassifierMkFitFit = _HLTInitialStepSequenceMkFitFit.copyAndExclude([HLTInitialStepHPSelectionSequence])
+_HLTInitialStepSequenceTrackCutClassifierMkFitFit += HLTInitialStepCutClassifierHPSelectionSequence 
+
+from Configuration.ProcessModifiers.trackCutClassifier_cff import trackCutClassifier
+(trackCutClassifier & ~ngtScouting).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackCutClassifier)
+(trackCutClassifier & trackingMkFitFit & ~ngtScouting).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackCutClassifierMkFitFit)

HLTrigger/Configuration/python/HLT_75e33/services/PyTorchService_cfi.py

@@ -0,0 +1,3 @@
+import FWCore.ParameterSet.Config as cms
+
+PyTorchService = cms.Service("PyTorchService")

HLTrigger/Configuration/python/HLT_75e33_cff.py

@@ -280,6 +280,7 @@
 fragment.load("HLTrigger/Configuration/HLT_75e33/services/FastTimerService_cfi")
 fragment.load("HLTrigger/Configuration/HLT_75e33/services/MessageLogger_cfi")
 fragment.load("HLTrigger/Configuration/HLT_75e33/services/ThroughputService_cfi")
+fragment.load("HLTrigger/Configuration/HLT_75e33/services/PyTorchService_cfi")
 
 fragment.schedule = cms.Schedule(*[
 

HLTrigger/Configuration/python/HLT_75e33_timing_cff.py

@@ -252,6 +252,7 @@
 fragment.load("HLTrigger/Configuration/HLT_75e33/sequences/HLTVertexRecoSequence_cfi")
 fragment.load("HLTrigger/Configuration/HLT_75e33/services/FastTimerService_cfi")
 fragment.load("HLTrigger/Configuration/HLT_75e33/services/ThroughputService_cfi")
+fragment.load("HLTrigger/Configuration/HLT_75e33/services/PyTorchService_cfi")
 
 fragment.schedule = cms.Schedule(*[
 

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"/>
+<!-- src/alpaka models the behavior of data formats -->
+<flags ALPAKA_BACKENDS="!serial"/>
 <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

@@ -38,3 +38,17 @@
 <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/PyTorch"/>
+  <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,127 @@
+#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"
+#include "TrackingTools/PatternTools/interface/Trajectory.h"
+#include "TrackingTools/PatternTools/interface/TrajTrackAssociation.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_;
+  edm::EDGetTokenT<std::vector<Trajectory>> trajectories_token_;
+  const edm::EDGetTokenT<PortableHostCollection<TrackTorchClassifierScoresSoA>> scores_token_;
+  const edm::EDGetTokenT<PortableHostCollection<TrackTorchClassifierFeaturesSoA>> features_token_;
+  const bool copy_trajectories_;
+  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<Trajectory>> filtered_trajectories_token_;
+  const edm::EDPutTokenT<TrajTrackAssociationCollection> traj_track_associations_token_;
+  const edm::EDPutTokenT<std::vector<float>> scores_output_token_;
+};
+
+TrackTorchClassifierFromSoA::TrackTorchClassifierFromSoA(const edm::ParameterSet& iConfig)
+    : tracks_token_(consumes(iConfig.getParameter<edm::InputTag>("src"))),
+      trajectories_token_(),
+      scores_token_(consumes(iConfig.getParameter<edm::InputTag>("scores"))),
+      features_token_(consumes(iConfig.getParameter<edm::InputTag>("features"))),
+      copy_trajectories_(iConfig.getParameter<bool>("copyTrajectories")),
+      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>()),
+      filtered_trajectories_token_(produces<std::vector<Trajectory>>()),
+      traj_track_associations_token_(produces<TrajTrackAssociationCollection>()),
+      scores_output_token_(produces<std::vector<float>>("MVAScores")) {
+  if (copy_trajectories_) {
+    trajectories_token_ = consumes(iConfig.getParameter<edm::InputTag>("src"));
+  }
+}
+
+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<bool>("copyTrajectories", false)
+      ->setComment("Whether to produce the filtered trajectory collection and associations");
+  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 std::vector<Trajectory>* trajectories = nullptr;
+  if (copy_trajectories_) {
+    trajectories = &iEvent.get(trajectories_token_);
+  }
+
+  const auto nTracks = tracks.size();
+
+  // Create filtered track collection and optionally copy the corresponding trajectories
+  auto filtered_tracks = std::make_unique<reco::TrackCollection>();
+  std::unique_ptr<std::vector<Trajectory>> filtered_trajectories;
+  std::unique_ptr<TrajTrackAssociationCollection> traj_track_associations;
+  if (copy_trajectories_) {
+    filtered_trajectories = std::make_unique<std::vector<Trajectory>>();
+    traj_track_associations = std::make_unique<TrajTrackAssociationCollection>(
+        iEvent.getRefBeforePut<std::vector<Trajectory>>(), iEvent.getRefBeforePut<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]);
+      if (copy_trajectories_) {
+        filtered_trajectories->push_back((*trajectories)[i]);
+        traj_track_associations->insert(
+            edm::Ref<std::vector<Trajectory>>(iEvent.getRefBeforePut<std::vector<Trajectory>>(),
+                                              filtered_trajectories->size() - 1),
+            reco::TrackRef(iEvent.getRefBeforePut<reco::TrackCollection>(), filtered_tracks->size() - 1));
+      }
+    }
+  }
+
+  iEvent.put(filtered_tracks_token_, std::move(filtered_tracks));
+  if (copy_trajectories_) {
+    iEvent.put(filtered_trajectories_token_, std::move(filtered_trajectories));
+    iEvent.put(traj_track_associations_token_, std::move(traj_track_associations));
+  }
+  iEvent.put(scores_output_token_, std::move(all_scores));
+}
+
+DEFINE_FWK_MODULE(TrackTorchClassifierFromSoA);