Fix energy regression issue in TICLv5

Moanwar · Moanwar · commit 57b365a35f6a · 2025-11-07T14:42:12.000+01:00
diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltTiclCandidate_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltTiclCandidate_cfi.py
@@ -6,14 +6,14 @@
     pluginInferenceAlgoTracksterInferenceByPFN = cms.PSet(
       algo_verbosity = cms.int32(0),
       onnxPIDModelPath = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/id_v0.onnx'),
-      onnxEnergyModelPath = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/energy_v0.onnx'),
+      onnxEnergyModelPath = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/energy_v1.onnx'),
       inputNames = cms.vstring(
         'input',
         'input_tr_features'
       ),
       output_en = cms.vstring('enreg_output'),
       output_id = cms.vstring('pid_output'),
-      eid_min_cluster_energy = cms.double(1),
+      eid_min_cluster_energy = cms.double(2.5),
       eid_n_layers = cms.int32(50),
       eid_n_clusters = cms.int32(10),
       doPID = cms.int32(1),
diff --git a/RecoHGCal/TICL/interface/TracksterInferenceAlgoBase.h b/RecoHGCal/TICL/interface/TracksterInferenceAlgoBase.h
@@ -18,6 +18,7 @@
 #include "FWCore/Framework/interface/ConsumesCollector.h"
 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
 #include "FWCore/PluginManager/interface/PluginFactory.h"
+#include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
 
 namespace ticl {
   class TracksterInferenceAlgoBase {
@@ -26,7 +27,9 @@ namespace ticl {
         : algo_verbosity_(conf.getParameter<int>("algo_verbosity")) {}
     virtual ~TracksterInferenceAlgoBase() {}
 
-    virtual void inputData(const std::vector<reco::CaloCluster>& layerClusters, std::vector<Trackster>& tracksters) = 0;
+    virtual void inputData(const std::vector<reco::CaloCluster>& layerClusters,
+                           std::vector<Trackster>& tracksters,
+                           const hgcal::RecHitTools& rhtools) = 0;
     virtual void runInference(std::vector<Trackster>& tracksters) = 0;
     static void fillPSetDescription(edm::ParameterSetDescription& desc) { desc.add<int>("algo_verbosity", 0); };
 
diff --git a/RecoHGCal/TICL/interface/TracksterInferenceByANN.h b/RecoHGCal/TICL/interface/TracksterInferenceByANN.h
@@ -7,7 +7,9 @@ namespace ticl {
   class TracksterInferenceByANN : public TracksterInferenceAlgoBase {
   public:
     explicit TracksterInferenceByANN(const edm::ParameterSet& conf);
-    void inputData(const std::vector<reco::CaloCluster>& layerClusters, std::vector<Trackster>& tracksters) override;
+    void inputData(const std::vector<reco::CaloCluster>& layerClusters,
+                   std::vector<Trackster>& tracksters,
+                   const hgcal::RecHitTools& rhtools) override;
     void runInference(std::vector<Trackster>& tracksters) override;
 
   private:
diff --git a/RecoHGCal/TICL/interface/TracksterInferenceByCNNv4.h b/RecoHGCal/TICL/interface/TracksterInferenceByCNNv4.h
@@ -9,7 +9,9 @@ namespace ticl {
   class TracksterInferenceByCNNv4 : public TracksterInferenceAlgoBase {
   public:
     explicit TracksterInferenceByCNNv4(const edm::ParameterSet& conf);
-    void inputData(const std::vector<reco::CaloCluster>& layerClusters, std::vector<Trackster>& tracksters) override;
+    void inputData(const std::vector<reco::CaloCluster>& layerClusters,
+                   std::vector<Trackster>& tracksters,
+                   const hgcal::RecHitTools& rhtools) override;
     void runInference(std::vector<Trackster>& tracksters) override;
 
     static void fillPSetDescription(edm::ParameterSetDescription& iDesc);
@@ -27,7 +29,6 @@ namespace ticl {
     int doPID_;
     int doRegression_;
 
-    hgcal::RecHitTools rhtools_;
     std::vector<std::vector<int64_t>> input_shapes_;
     std::vector<int> tracksterIndices_;
     std::vector<std::vector<float>> input_Data_;
diff --git a/RecoHGCal/TICL/interface/TracksterInferenceByDNN.h b/RecoHGCal/TICL/interface/TracksterInferenceByDNN.h
@@ -9,7 +9,9 @@ namespace ticl {
   class TracksterInferenceByDNN : public TracksterInferenceAlgoBase {
   public:
     explicit TracksterInferenceByDNN(const edm::ParameterSet& conf);
-    void inputData(const std::vector<reco::CaloCluster>& layerClusters, std::vector<Trackster>& tracksters) override;
+    void inputData(const std::vector<reco::CaloCluster>& layerClusters,
+                   std::vector<Trackster>& tracksters,
+                   const hgcal::RecHitTools& rhtools) override;
     void runInference(std::vector<Trackster>& tracksters) override;
 
     static void fillPSetDescription(edm::ParameterSetDescription& iDesc);
@@ -30,7 +32,6 @@ namespace ticl {
     int doPID_;
     int doRegression_;
 
-    hgcal::RecHitTools rhtools_;
     std::vector<std::vector<int64_t>> input_shapes_;
     std::vector<int> tracksterIndices_;
     std::vector<std::vector<float>> input_Data_;
diff --git a/RecoHGCal/TICL/interface/TracksterInferenceByPFN.h b/RecoHGCal/TICL/interface/TracksterInferenceByPFN.h
@@ -9,7 +9,9 @@ namespace ticl {
   class TracksterInferenceByPFN : public TracksterInferenceAlgoBase {
   public:
     explicit TracksterInferenceByPFN(const edm::ParameterSet& conf);
-    void inputData(const std::vector<reco::CaloCluster>& layerClusters, std::vector<Trackster>& tracksters) override;
+    void inputData(const std::vector<reco::CaloCluster>& layerClusters,
+                   std::vector<Trackster>& tracksters,
+                   const hgcal::RecHitTools& rhtools) override;
     void runInference(std::vector<Trackster>& tracksters) override;
 
     static void fillPSetDescription(edm::ParameterSetDescription& iDesc);
@@ -30,7 +32,6 @@ namespace ticl {
     int doPID_;
     int doRegression_;
 
-    hgcal::RecHitTools rhtools_;
     std::vector<std::vector<int64_t>> input_shapes_;
     std::vector<int> tracksterIndices_;
     std::vector<std::vector<float>> input_Data_;
diff --git a/RecoHGCal/TICL/plugins/TICLCandidateProducer.cc b/RecoHGCal/TICL/plugins/TICLCandidateProducer.cc
@@ -332,7 +332,7 @@ void TICLCandidateProducer::produce(edm::Event &evt, const edm::EventSetup &es)
                         true);
   if (regressionAndPid_) {
     // Run inference algorithm
-    inferenceAlgo_->inputData(layerClusters, *resultTracksters);
+    inferenceAlgo_->inputData(layerClusters, *resultTracksters, rhtools_);
     inferenceAlgo_->runInference(
         *resultTracksters);  //option to use "Linking" instead of "CLU3D"/"energyAndPid" instead of "PID"
   }
diff --git a/RecoHGCal/TICL/plugins/TracksterInferenceByANN.cc b/RecoHGCal/TICL/plugins/TracksterInferenceByANN.cc
@@ -10,7 +10,8 @@ namespace ticl {
   }
 
   void TracksterInferenceByANN::inputData(const std::vector<reco::CaloCluster>& layerClusters,
-                                          std::vector<Trackster>& tracksters) {
+                                          std::vector<Trackster>& tracksters,
+                                          const hgcal::RecHitTools& rhtools) {
     // Prepare data for inference
   }
 
diff --git a/RecoHGCal/TICL/plugins/TracksterInferenceByCNNv4.cc b/RecoHGCal/TICL/plugins/TracksterInferenceByCNNv4.cc
@@ -30,7 +30,8 @@ namespace ticl {
 
   // Method to process input data and prepare it for inference
   void TracksterInferenceByCNNv4::inputData(const std::vector<reco::CaloCluster>& layerClusters,
-                                            std::vector<Trackster>& tracksters) {
+                                            std::vector<Trackster>& tracksters,
+                                            const hgcal::RecHitTools& rhtools) {
     tracksterIndices_.clear();  // Clear previous indices
     for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
       float sumClusterEnergy = 0.;
@@ -74,7 +75,7 @@ namespace ticl {
       // Fill input data with cluster information
       for (const int& k : clusterIndices) {
         const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
-        int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
+        int j = rhtools.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
         if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
           auto index = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
           input_Data_[0][index] =
diff --git a/RecoHGCal/TICL/plugins/TracksterInferenceByDNN.cc b/RecoHGCal/TICL/plugins/TracksterInferenceByDNN.cc
@@ -33,12 +33,13 @@ namespace ticl {
 
   // Method to process input data and prepare it for inference
   void TracksterInferenceByDNN::inputData(const std::vector<reco::CaloCluster>& layerClusters,
-                                          std::vector<Trackster>& tracksters) {
+                                          std::vector<Trackster>& tracksters,
+                                          const hgcal::RecHitTools& rhtools) {
     tracksterIndices_.clear();  // Clear previous indices
     for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
       float sumClusterEnergy = 0.;
       for (const unsigned int& vertex : tracksters[i].vertices()) {
-        if (rhtools_.isBarrel(layerClusters[vertex].seed()))
+        if (rhtools.isBarrel(layerClusters[vertex].seed()))
           continue;
         sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
         if (sumClusterEnergy >= eidMinClusterEnergy_) {
@@ -79,7 +80,7 @@ namespace ticl {
       // Fill input data with cluster information
       for (const int& k : clusterIndices) {
         const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
-        int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
+        int j = rhtools.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
         if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
           auto index = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
           input_Data_[0][index] =
diff --git a/RecoHGCal/TICL/plugins/TracksterInferenceByPFN.cc b/RecoHGCal/TICL/plugins/TracksterInferenceByPFN.cc
@@ -29,24 +29,19 @@ namespace ticl {
     onnxPIDSession_ = onnxPIDRuntimeInstance_.get();
     onnxEnergySession_ = onnxEnergyRuntimeInstance_.get();
   }
-
   // Method to process input data and prepare it for inference
   void TracksterInferenceByPFN::inputData(const std::vector<reco::CaloCluster>& layerClusters,
-                                          std::vector<Trackster>& tracksters) {
+                                          std::vector<Trackster>& tracksters,
+                                          const hgcal::RecHitTools& rhtools) {
     tracksterIndices_.clear();  // Clear previous indices
     for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
-      float sumClusterEnergy = 0.;
       for (const unsigned int& vertex : tracksters[i].vertices()) {
-        if (rhtools_.isBarrel(layerClusters[vertex].seed()))
+        if (rhtools.isBarrel(layerClusters[vertex].seed()))
           continue;
-        sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
-        if (sumClusterEnergy >= eidMinClusterEnergy_) {
-          tracksters[i].setRegressedEnergy(0.f);  // Set regressed energy to 0
-          tracksters[i].zeroProbabilities();      // Zero out probabilities
-          tracksterIndices_.push_back(i);         // Add index to the list
-          break;
-        }
       }
+      tracksters[i].setRegressedEnergy(0.f);  // Initialize regressed energy to 0
+      tracksters[i].zeroProbabilities();      // Zero out probabilities
+      tracksterIndices_.push_back(i);         // Only tracksters above threshold go to inference
     }
 
     // Prepare input shapes and data for inference
@@ -88,7 +83,7 @@ namespace ticl {
       // Fill input data with cluster information
       for (const int& k : clusterIndices) {
         const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
-        int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
+        int j = rhtools.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
         if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
           auto index_lc = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
           // Adding more features regarding LC, such as E, eta, phi, x, y, z, and nhits.
@@ -114,11 +109,12 @@ namespace ticl {
     if (doPID_ and doRegression_) {
       // Run energy model inference
       auto result = onnxEnergySession_->run(inputNames_, input_Data_, input_shapes_, output_en_, batchSize_);
-      auto& energyOutputTensor = result[0];
+      auto& energyOutput = result[0];
       if (!output_en_.empty()) {
         for (int i = 0; i < static_cast<int>(batchSize_); i++) {
-          const float energy = energyOutputTensor[i];
-          tracksters[tracksterIndices_[i]].setRegressedEnergy(energy);  // Update energy
+          auto& ts = tracksters[tracksterIndices_[i]];
+          float energy = ts.raw_energy() > eidMinClusterEnergy_ ? energyOutput[i] : ts.raw_energy();
+          ts.setRegressedEnergy(energy);
         }
       }
     }
diff --git a/RecoHGCal/TICL/plugins/TracksterLinksProducer.cc b/RecoHGCal/TICL/plugins/TracksterLinksProducer.cc
@@ -239,7 +239,7 @@ void TracksterLinksProducer::produce(edm::Event &evt, const edm::EventSetup &es)
 
   if (regressionAndPid_) {
     // Run inference algorithm
-    inferenceAlgo_->inputData(layerClusters, *resultTracksters);
+    inferenceAlgo_->inputData(layerClusters, *resultTracksters, rhtools_);
     inferenceAlgo_->runInference(
         *resultTracksters);  //option to use "Linking" instead of "CLU3D"/"energyAndPid" instead of "PID"
   }
diff --git a/RecoHGCal/TICL/plugins/TrackstersProducer.cc b/RecoHGCal/TICL/plugins/TrackstersProducer.cc
@@ -28,6 +28,11 @@
 #include "FWCore/Framework/interface/ConsumesCollector.h"
 #include "RecoHGCal/TICL/interface/TracksterInferenceAlgoFactory.h"
 
+#include "Geometry/HGCalCommonData/interface/HGCalDDDConstants.h"
+#include "Geometry/Records/interface/IdealGeometryRecord.h"
+#include "Geometry/CommonDetUnit/interface/GeomDet.h"
+#include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
+
 using namespace ticl;
 using namespace cms::Ort;
 
@@ -42,6 +47,10 @@ class TrackstersProducer : public edm::stream::EDProducer<> {
   // static methods for handling the global cache
   static std::unique_ptr<TrackstersCache> initializeGlobalCache(const edm::ParameterSet&);
   static void globalEndJob(TrackstersCache*);
+  void beginRun(const edm::Run&, const edm::EventSetup& es) override {
+    const auto& geom = es.getData(geometry_token_);
+    rhtools_.setGeometry(geom);
+  }
 
 private:
   std::string detector_;
@@ -56,8 +65,10 @@ class TrackstersProducer : public edm::stream::EDProducer<> {
   edm::EDGetTokenT<TICLLayerTiles> layer_clusters_tiles_token_;
   edm::EDGetTokenT<TICLLayerTilesHFNose> layer_clusters_tiles_hfnose_token_;
   const edm::EDGetTokenT<std::vector<TICLSeedingRegion>> seeding_regions_token_;
-  const std::string itername_;
   ticl::Trackster::IterationIndex iterIndex_ = ticl::Trackster::IterationIndex(0);
+  const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geometry_token_;
+  hgcal::RecHitTools rhtools_;
+  const std::string itername_;
 };
 DEFINE_FWK_MODULE(TrackstersProducer);
 
@@ -71,6 +82,7 @@ TrackstersProducer::TrackstersProducer(const edm::ParameterSet& ps)
           consumes<edm::ValueMap<std::pair<float, float>>>(ps.getParameter<edm::InputTag>("time_layerclusters"))),
       seeding_regions_token_(
           consumes<std::vector<TICLSeedingRegion>>(ps.getParameter<edm::InputTag>("seeding_regions"))),
+      geometry_token_(esConsumes<CaloGeometry, CaloGeometryRecord, edm::Transition::BeginRun>()),
       itername_(ps.getParameter<std::string>("itername")) {
   auto plugin = ps.getParameter<std::string>("patternRecognitionBy");
   auto pluginPSet = ps.getParameter<edm::ParameterSet>("pluginPatternRecognitionBy" + plugin);
@@ -192,7 +204,7 @@ void TrackstersProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
 
       myAlgoHFNose_->makeTracksters(inputHFNose, *initialResult, seedToTrackstersAssociation);
       // Run inference algorithm
-      inferenceAlgo_->inputData(layerClusters, *initialResult);
+      inferenceAlgo_->inputData(layerClusters, *initialResult, rhtools_);
       inferenceAlgo_->runInference(*initialResult);
       myAlgoHFNose_->filter(*result, *initialResult, inputHFNose, seedToTrackstersAssociation);
 
@@ -203,7 +215,7 @@ void TrackstersProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
 
       myAlgo_->makeTracksters(input, *initialResult, seedToTrackstersAssociation);
       // Run inference algorithm
-      inferenceAlgo_->inputData(layerClusters, *initialResult);
+      inferenceAlgo_->inputData(layerClusters, *initialResult, rhtools_);
       inferenceAlgo_->runInference(*initialResult);
       myAlgo_->filter(*result, *initialResult, input, seedToTrackstersAssociation);
     }
diff --git a/RecoHGCal/TICL/python/iterativeTICL_cff.py b/RecoHGCal/TICL/python/iterativeTICL_cff.py
@@ -107,10 +107,10 @@
         inputNames  = cms.vstring('input','input_tr_features'),
         output_en   = cms.vstring('enreg_output'),
         output_id   = cms.vstring('pid_output'),
-	eid_min_cluster_energy = cms.double(1),
+	eid_min_cluster_energy = cms.double(2.5),
         eid_n_clusters = cms.int32(10),
 	eid_n_layers = cms.int32(50),
-        onnxEnergyModelPath = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/energy_v0.onnx'),
+        onnxEnergyModelPath = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/energy_v1.onnx'),
         onnxPIDModelPath = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/id_v0.onnx'),
         type = cms.string('TracksterInferenceByPFN')
     )
@@ -124,14 +124,14 @@
         onnxPIDModelPath=cms.FileInPath(
             'RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/id_v0.onnx'),
         onnxEnergyModelPath=cms.FileInPath(
-            'RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/energy_v0.onnx'),
+            'RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/linking/energy_v1.onnx'),
         inputNames=cms.vstring(
             'input',
             'input_tr_features'
         ),
         output_en=cms.vstring('enreg_output'),
         output_id=cms.vstring('pid_output'),
-        eid_min_cluster_energy=cms.double(1),
+        eid_min_cluster_energy=cms.double(2.5),
         eid_n_layers=cms.int32(50),
         eid_n_clusters=cms.int32(10),
         doPID=cms.int32(1),
diff --git a/RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h b/RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h
@@ -112,6 +112,7 @@ namespace hgcal {
 
   private:
     const CaloGeometry* geom_;
+    void checkGeometry() const;
     unsigned int eeOffset_, fhOffset_, bhFirstLayer_, bhLastLayer_, bhOffset_, fhLastLayer_, noseLastLayer_;
     unsigned int hcalBarrelFirstLayer_, hcalBarrelLastLayer_, ecalBarrelFirstLayer_, ecalBarrelLastLayer_;
     unsigned int maxNumberOfWafersPerLayer_, maxNumberOfWafersNose_;
diff --git a/RecoLocalCalo/HGCalRecAlgos/src/RecHitTools.cc b/RecoLocalCalo/HGCalRecAlgos/src/RecHitTools.cc

Original file line number	Diff line number	Diff line change
`@@ -332,7 +332,7 @@ void TICLCandidateProducer::produce(edm::Event &evt, const edm::EventSetup &es)`
`332`	`332`	`true);`
`333`	`333`	`if (regressionAndPid_) {`
`334`	`334`	`// Run inference algorithm`
`335`		`- inferenceAlgo_->inputData(layerClusters, *resultTracksters);`
	`335`	`+ inferenceAlgo_->inputData(layerClusters, *resultTracksters, rhtools_);`
`336`	`336`	`inferenceAlgo_->runInference(`
`337`	`337`	`*resultTracksters); //option to use "Linking" instead of "CLU3D"/"energyAndPid" instead of "PID"`
`338`	`338`	`}`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,8 @@ namespace ticl {`
`10`	`10`	`}`
`11`	`11`
`12`	`12`	`void TracksterInferenceByANN::inputData(const std::vector<reco::CaloCluster>& layerClusters,`
`13`		`- std::vector<Trackster>& tracksters) {`
	`13`	`+ std::vector<Trackster>& tracksters,`
	`14`	`+ const hgcal::RecHitTools& rhtools) {`
`14`	`15`	`// Prepare data for inference`
`15`	`16`	`}`
`16`	`17`