DNN v3 with timing info, corrected seed pT

gsokmen · gsokmen · commit ad12e5250296 · 2025-08-26T19:00:27.000+02:00
diff --git a/RecoHGCal/TICL/interface/SuperclusteringDNNInputs.h b/RecoHGCal/TICL/interface/SuperclusteringDNNInputs.h
@@ -2,6 +2,10 @@
 // Author: Theo Cuisset - theo.cuisset@cern.ch
 // Date: 11/2023
 
+// Modified by Gamze Sokmen 
+// Changes: Implementation of the delta time feature under a new DNN input version (v3) for the superclustering DNN and correcting the seed pT calculation.
+// Date: 07/2025
+
 #ifndef __RecoHGCal_TICL_SuperclusteringDNNInputs_H__
 #define __RecoHGCal_TICL_SuperclusteringDNNInputs_H__
 
@@ -87,6 +91,38 @@ namespace ticl {
     }
   };
 
+  /* Third version of DNN by Gamze Sokmen, making use of time information as new variables.
+  Uses features : ['DeltaEta', 'DeltaPhi', 'multi_en', 'multi_eta', 'multi_pt', 'seedEta','seedPhi','seedEn', 'seedPt',   theta', 'theta_xz_seedFrame', 'theta_yz_seedFrame', 'theta_xy_cmsFrame', 'theta_yz_cmsFrame', 'theta_xz_cmsFrame', 'explVar', 'explVarRatio', 'mod_deltaTime']
+  */
+ 
+  class SuperclusteringDNNInputV3 : public AbstractSuperclusteringDNNInput {
+    public:
+      unsigned int featureCount() const override { return 18; }
+
+      std::vector<float> computeVector(ticl::Trackster const& ts_base, ticl::Trackster const& ts_toCluster) override;
+
+      std::vector<std::string> featureNames() const override {
+        return {"DeltaEtaBaryc",
+                "DeltaPhiBaryc",
+                "multi_en",
+                "multi_eta",
+                "multi_pt",
+                "seedEta",
+                "seedPhi",
+                "seedEn",
+                "seedPt",
+                "theta",
+                "theta_xz_seedFrame",
+                "theta_yz_seedFrame",
+                "theta_xy_cmsFrame",
+                "theta_yz_cmsFrame",
+                "theta_xz_cmsFrame",
+                "explVar",
+                "explVarRatio",
+                "mod_deltaTime"};
+      }
+    };  
+
   std::unique_ptr<AbstractSuperclusteringDNNInput> makeSuperclusteringDNNInputFromString(std::string dnnVersion);
 }  // namespace ticl
 
diff --git a/RecoHGCal/TICL/plugins/SuperclusteringSampleDumper.cc b/RecoHGCal/TICL/plugins/SuperclusteringSampleDumper.cc
@@ -151,7 +151,7 @@ void SuperclusteringSampleDumper::analyze(const edm::Event& evt, const edm::Even
   std::iota(trackstersIndicesPt.begin(), trackstersIndicesPt.end(), 0);
   std::stable_sort(
       trackstersIndicesPt.begin(), trackstersIndicesPt.end(), [&inputTracksters](unsigned int i1, unsigned int i2) {
-        return (*inputTracksters)[i1].raw_pt() > (*inputTracksters)[i2].raw_pt();
+	return (*inputTracksters)[i1].raw_energy()*std::sin((*inputTracksters)[i1].barycenter().Theta()) > (*inputTracksters)[i2].raw_energy()*std::sin((*inputTracksters)[i2].barycenter().Theta());
       });
 
   // Order of loops are reversed compared to SuperclusteringProducer (here outer is seed, inner is candidate), for performance reasons.
@@ -162,7 +162,7 @@ void SuperclusteringSampleDumper::analyze(const edm::Event& evt, const edm::Even
         trackstersIndicesPt[ts_seed_idx_pt];  // Index of seed trackster in input collection (not in pT sorted collection)
     Trackster const& ts_seed = (*inputTracksters)[ts_seed_idx_input];
 
-    if (ts_seed.raw_pt() < seedPtThreshold_)
+    if (ts_seed.raw_energy()*std::sin(ts_seed.barycenter().Theta()) < seedPtThreshold_)
       break;  // All further seeds will have lower pT than threshold (due to pT sorting)
 
     if (!checkExplainedVarianceRatioCut(ts_seed))
@@ -265,8 +265,8 @@ void SuperclusteringSampleDumper::fillDescriptions(edm::ConfigurationDescription
       ->setComment("Input trackster collection, same as what is used for superclustering inference.");
   desc.add<edm::InputTag>("recoToSimAssociatorCP",
                           edm::InputTag("tracksterSimTracksterAssociationLinkingbyCLUE3D", "recoToSim"));
-  desc.ifValue(edm::ParameterDescription<std::string>("dnnInputsVersion", "v2", true),
-               edm::allowedValues<std::string>("v1", "v2"))
+  desc.ifValue(edm::ParameterDescription<std::string>("dnnInputsVersion", "v3", true),
+               edm::allowedValues<std::string>("v1", "v2", "v3"))
       ->setComment(
           "DNN inputs version tag. Defines which set of features is fed to the DNN. Must match with the actual DNN.");
   // Cuts are intentionally looser than those used for inference in TracksterLinkingBySuperClustering.cpp
diff --git a/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClusteringDNN.cc b/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClusteringDNN.cc
@@ -116,7 +116,7 @@ void TracksterLinkingbySuperClusteringDNN::linkTracksters(
   std::iota(trackstersIndicesPt.begin(), trackstersIndicesPt.end(), 0);
   std::stable_sort(
       trackstersIndicesPt.begin(), trackstersIndicesPt.end(), [&inputTracksters](unsigned int i1, unsigned int i2) {
-        return inputTracksters[i1].raw_pt() > inputTracksters[i2].raw_pt();
+        return inputTracksters[i1].raw_energy()*std::sin(inputTracksters[i1].barycenter().Theta()) > inputTracksters[i2].raw_energy()*std::sin(inputTracksters[i2].barycenter().Theta());
       });
 
   /* Evaluate in minibatches since running with trackster count = 3000 leads to a short-lived ~15GB memory allocation
@@ -146,7 +146,8 @@ void TracksterLinkingbySuperClusteringDNN::linkTracksters(
     Trackster const& ts_cand = inputTracksters[trackstersIndicesPt[ts_cand_idx_pt]];
 
     if (ts_cand.raw_energy() < candidateEnergyThreshold_ ||
-        !checkExplainedVarianceRatioCut(ts_cand))  // || !trackstersPassesPIDCut(ts_cand)
+	//        !checkExplainedVarianceRatioCut(ts_cand))  // || !trackstersPassesPIDCut(ts_cand)
+	!checkExplainedVarianceRatioCut(ts_cand))//   || !trackstersPassesPIDCut(ts_cand))
       continue;
 
     auto& tracksterTiles = tracksterTilesBothEndcaps_pt[ts_cand.barycenter().eta() > 0];
@@ -164,7 +165,7 @@ void TracksterLinkingbySuperClusteringDNN::linkTracksters(
 
           Trackster const& ts_seed = inputTracksters[trackstersIndicesPt[ts_seed_idx_pt]];
 
-          if (ts_seed.raw_pt() < seedPtThreshold_)
+          if (ts_seed.raw_energy()*std::sin(ts_seed.barycenter().Theta()) < seedPtThreshold_)
             break;  // All further seeds will have lower pT than threshold (due to pT sorting)
 
           if (!checkExplainedVarianceRatioCut(ts_seed) || !trackstersPassesPIDCut(ts_seed))
@@ -245,87 +246,95 @@ void TracksterLinkingbySuperClusteringDNN::linkTracksters(
   Also mask seeds (only needed to add tracksters not in a supercluster to the output). */
   std::vector<bool> tracksterMask(tracksterCount, false);
 
-  /* Index of the seed trackster of the previous iteration 
-  Initialized with an id that cannot be obtained in input */
+
+  /////////////////////////////////////////////////////////////////////////TRKBUILDINGMOD
+  
   unsigned int previousCandTrackster_idx = std::numeric_limits<unsigned int>::max();
   unsigned int bestSeedForCurrentCandidate_idx = std::numeric_limits<unsigned int>::max();
   float bestSeedForCurrentCandidate_dnnScore = nnWorkingPoint_;
-
-  // Lambda to be called when there is a transition from one candidate to the next (as well as after the last iteration)
-  // Does the actual supercluster creation
+  
+  // Track which tracksters were ever used as candidates
+  std::vector<bool> usedAsCandidate(tracksterCount, false);
+  
+  
   auto onCandidateTransition = [&](unsigned ts_cand_idx) {
-    if (bestSeedForCurrentCandidate_idx <
-        std::numeric_limits<unsigned int>::max()) {  // At least one seed can be superclustered with the candidate
-      tracksterMask[ts_cand_idx] = true;  // Mask the candidate so it is not considered as seed in later iterations
-
-      // Look for a supercluster of the seed
-      std::vector<std::vector<unsigned int>>::iterator seed_supercluster_it =
-          std::find_if(outputSuperclusters.begin(),
-                       outputSuperclusters.end(),
-                       [bestSeedForCurrentCandidate_idx](std::vector<unsigned int> const& sc) {
-                         return sc[0] == bestSeedForCurrentCandidate_idx;
-                       });
-
-      if (seed_supercluster_it == outputSuperclusters.end()) {  // No supercluster exists yet for the seed. Create one.
-        outputSuperclusters.emplace_back(std::initializer_list<unsigned int>{bestSeedForCurrentCandidate_idx});
-        resultTracksters.emplace_back(inputTracksters[bestSeedForCurrentCandidate_idx]);
-        linkedTracksterIdToInputTracksterId.emplace_back(
-            std::initializer_list<unsigned int>{bestSeedForCurrentCandidate_idx});
-        seed_supercluster_it = outputSuperclusters.end() - 1;
-        tracksterMask[bestSeedForCurrentCandidate_idx] =
-            true;  // mask the seed as well (needed to find tracksters not in any supercluster)
+    if (bestSeedForCurrentCandidate_idx < std::numeric_limits<unsigned int>::max()) {
+      tracksterMask[ts_cand_idx] = true;           // Mask the candidate so it’s not reused as a seed
+      usedAsCandidate[ts_cand_idx] = true;
+      
+      // Find the supercluster the seed belongs to (even if it's already used in another supercluster)
+      // Find existing supercluster for the seed
+      auto seed_supercluster_it =
+        std::find_if(outputSuperclusters.begin(),
+                     outputSuperclusters.end(),
+                     [bestSeedForCurrentCandidate_idx](const std::vector<unsigned int>& sc) {
+                       return sc[0] == bestSeedForCurrentCandidate_idx;
+                     });
+      if (seed_supercluster_it == outputSuperclusters.end()) {
+	// No supercluster exists for this seed, create one
+	outputSuperclusters.emplace_back(std::initializer_list<unsigned int>{bestSeedForCurrentCandidate_idx});
+	resultTracksters.emplace_back(inputTracksters[bestSeedForCurrentCandidate_idx]);
+	linkedTracksterIdToInputTracksterId.emplace_back(
+							 std::initializer_list<unsigned int>{bestSeedForCurrentCandidate_idx});
+	seed_supercluster_it = outputSuperclusters.end() - 1;
+	tracksterMask[bestSeedForCurrentCandidate_idx] = true;
       }
-      // Index of the supercluster into resultTracksters, outputSuperclusters and linkedTracksterIdToInputTracksterId collections (the indices are the same)
+      
       unsigned int indexIntoOutputTracksters = seed_supercluster_it - outputSuperclusters.begin();
       seed_supercluster_it->push_back(ts_cand_idx);
       resultTracksters[indexIntoOutputTracksters].mergeTracksters(inputTracksters[ts_cand_idx]);
       linkedTracksterIdToInputTracksterId[indexIntoOutputTracksters].push_back(ts_cand_idx);
-
+      
       assert(outputSuperclusters.size() == resultTracksters.size() &&
-             outputSuperclusters.size() == linkedTracksterIdToInputTracksterId.size());
+	     outputSuperclusters.size() == linkedTracksterIdToInputTracksterId.size());
       assert(seed_supercluster_it->size() == linkedTracksterIdToInputTracksterId[indexIntoOutputTracksters].size());
-
+      
       bestSeedForCurrentCandidate_idx = std::numeric_limits<unsigned int>::max();
       bestSeedForCurrentCandidate_dnnScore = nnWorkingPoint_;
     }
   };
-
-  //Iterate over minibatches
+  
+  // Iterate over minibatches
   for (unsigned int batchIndex = 0; batchIndex < batchOutputs.size(); batchIndex++) {
-    std::vector<float> const& currentBatchOutputs = batchOutputs[batchIndex];  // DNN score outputs
-    // Iterate over seed-candidate pairs inside current minibatch
+    std::vector<float> const& currentBatchOutputs = batchOutputs[batchIndex];
+    
     for (unsigned int indexInBatch = 0; indexInBatch < tracksterIndicesUsedInDNN[batchIndex].size(); indexInBatch++) {
       assert(indexInBatch < static_cast<unsigned int>(batchOutputs[batchIndex].size()));
-
+      
       const unsigned int ts_seed_idx = tracksterIndicesUsedInDNN[batchIndex][indexInBatch].first;
       const unsigned int ts_cand_idx = tracksterIndicesUsedInDNN[batchIndex][indexInBatch].second;
       const float currentDnnScore = currentBatchOutputs[indexInBatch];
-
+      
       if (previousCandTrackster_idx != std::numeric_limits<unsigned int>::max() &&
-          ts_cand_idx != previousCandTrackster_idx) {
-        // There is a transition from one seed to the next (don't make a transition for the first iteration)
-        onCandidateTransition(previousCandTrackster_idx);
+	  ts_cand_idx != previousCandTrackster_idx) {
+	onCandidateTransition(previousCandTrackster_idx);
       }
-
-      if (currentDnnScore > bestSeedForCurrentCandidate_dnnScore && !tracksterMask[ts_seed_idx]) {
-        // Check that the DNN suggests superclustering, that this seed-candidate assoc is better than previous ones, and that the seed is not already in a supercluster as candidate
-        bestSeedForCurrentCandidate_idx = ts_seed_idx;
-        bestSeedForCurrentCandidate_dnnScore = currentDnnScore;
+      
+      // Ignore seed if it was previously used as a candidate
+      if (currentDnnScore > bestSeedForCurrentCandidate_dnnScore  && !usedAsCandidate[ts_seed_idx]) {
+	bestSeedForCurrentCandidate_idx = ts_seed_idx;
+	bestSeedForCurrentCandidate_dnnScore = currentDnnScore;
       }
+      
       previousCandTrackster_idx = ts_cand_idx;
     }
   }
   onCandidateTransition(previousCandTrackster_idx);
-
-  // Adding one-trackster superclusters for all tracksters not in a supercluster already that pass the seed threshold
+  
+  // Create singleton superclusters for unused tracksters with enough pt
   for (unsigned int ts_id = 0; ts_id < tracksterCount; ts_id++) {
-    if (!tracksterMask[ts_id] && inputTracksters[ts_id].raw_pt() >= seedPtThreshold_) {
+    if (!tracksterMask[ts_id] && inputTracksters[ts_id].raw_energy()*std::sin(inputTracksters[ts_id].barycenter().Theta()) >= seedPtThreshold_) {
       outputSuperclusters.emplace_back(std::initializer_list<unsigned int>{ts_id});
       resultTracksters.emplace_back(inputTracksters[ts_id]);
       linkedTracksterIdToInputTracksterId.emplace_back(std::initializer_list<unsigned int>{ts_id});
     }
   }
-
+  
+  
+  
+  /////////////////////////////////////////////////////////////////////////TRKBUILDINGMOD
+  
+  
 #ifdef EDM_ML_DEBUG
   for (std::vector<unsigned int> const& sc : outputSuperclusters) {
     std::ostringstream s;
@@ -340,8 +349,8 @@ void TracksterLinkingbySuperClusteringDNN::linkTracksters(
 void TracksterLinkingbySuperClusteringDNN::fillPSetDescription(edm::ParameterSetDescription& desc) {
   TracksterLinkingAlgoBase::fillPSetDescription(desc);  // adds algo_verbosity
   desc.add<edm::FileInPath>("onnxModelPath")->setComment("Path to DNN (as ONNX model)");
-  desc.ifValue(edm::ParameterDescription<std::string>("dnnInputsVersion", "v2", true),
-               edm::allowedValues<std::string>("v1", "v2"))
+  desc.ifValue(edm::ParameterDescription<std::string>("dnnInputsVersion", "v3", true),
+               edm::allowedValues<std::string>("v1", "v2", "v3"))
       ->setComment(
           "DNN inputs version tag. Defines which set of features is fed to the DNN. Must match with the actual DNN.");
   desc.add<unsigned int>("inferenceBatchSize", 1e5)
@@ -379,4 +388,4 @@ void TracksterLinkingbySuperClusteringDNN::fillPSetDescription(edm::ParameterSet
           {static_cast<int>(Trackster::ParticleType::photon), static_cast<int>(Trackster::ParticleType::electron)})
       ->setComment("List of PID particle types (ticl::Trackster::ParticleType enum) to consider for PID filtering");
   desc.add<double>("PIDThreshold", 0.8)->setComment("PID score threshold");
-}
+}
diff --git a/RecoHGCal/TICL/src/SuperclusteringDNNInputs.cc b/RecoHGCal/TICL/src/SuperclusteringDNNInputs.cc
