Electron superclustering in TICL using a DNN

CommonTools/RecoAlgos/interface/MultiVectorManager.h

@@ -53,6 +53,12 @@ class MultiVectorManager {
 
   class Iterator {
   public:
+    using iterator_category = std::forward_iterator_tag;
+    using difference_type = std::ptrdiff_t;
+    using value_type = T;
+    using pointer = T*;
+    using reference = T&;
+
     Iterator(const MultiVectorManager& manager, size_t index) : manager(manager), currentIndex(index) {}
 
     bool operator!=(const Iterator& other) const { return currentIndex != other.currentIndex; }

Configuration/ProcessModifiers/python/ticl_superclustering_dnn_cff.py

@@ -0,0 +1,10 @@
+import FWCore.ParameterSet.Config as cms
+
+from Configuration.ProcessModifiers.ticl_v5_cff import ticl_v5
+from Configuration.ProcessModifiers.ticl_superclustering_mustache_pf_cff import ticl_superclustering_mustache_pf
+from Configuration.ProcessModifiers.ticl_superclustering_mustache_ticl_cff import ticl_superclustering_mustache_ticl
+
+# Modifier to run superclustering with DNN
+# It is not to be used directly in a Process since it is currently the default for ticl_v5
+# It only exists as convenience in configuration files 
+ticl_superclustering_dnn = ticl_v5 & (~ticl_superclustering_mustache_pf) & (~ticl_superclustering_mustache_ticl)

Configuration/ProcessModifiers/python/ticl_superclustering_mustache_pf_cff.py

@@ -0,0 +1,4 @@
+import FWCore.ParameterSet.Config as cms
+
+# Modifier (to be applied on top of ticl_v5 modifier) to use the pre-TICLv5 superclustering code (that translates tracksters to PFClusters before running Mustache)
+ticl_superclustering_mustache_pf = cms.Modifier() 

Configuration/ProcessModifiers/python/ticl_superclustering_mustache_ticl_cff.py

@@ -0,0 +1,4 @@
+import FWCore.ParameterSet.Config as cms
+
+# Modifier (to be applied on top of ticl_v5 modifier) to run superclustering using Mustache inside TICL (with tracksters as inputs)
+ticl_superclustering_mustache_ticl = cms.Modifier() 

Configuration/PyReleaseValidation/README.md

@@ -96,6 +96,9 @@ The offsets currently in use are:
 * 0.103: Phase-2 aging, 3000fb-1
 * 0.201: HGCAL special TICL Pattern recognition Workflows: clue3D
 * 0.202: HGCAL special TICL Pattern recognition Workflows: FastJet
+* 0.203: HGCAL TICLv5
+* 0.204: HGCAL superclustering : using Mustache in TICLv5
+* 0.205: HGCAL superclustering : using old PFCluster-based Mustache algorithm with TICLv5
 * 0.302: FastSim Run-3 trackingOnly validation
 * 0.303: FastSim Run-3 MB for mixing
 * 0.9001: Sonic Triton

Configuration/PyReleaseValidation/python/upgradeWorkflowComponents.py

@@ -716,6 +716,44 @@ def condition(self, fragment, stepList, key, hasHarvest):
 upgradeWFs['ticl_v5'].step3 = {'--procModifiers': 'ticl_v5'}
 upgradeWFs['ticl_v5'].step4 = {'--procModifiers': 'ticl_v5'}
 
+class UpgradeWorkflow_ticl_v5_superclustering(UpgradeWorkflow):
+    def setup_(self, step, stepName, stepDict, k, properties):
+        if 'RecoGlobal' in step:
+            stepDict[stepName][k] = merge([self.step3, stepDict[step][k]])
+        if 'HARVESTGlobal' in step:
+            stepDict[stepName][k] = merge([self.step4, stepDict[step][k]])
+    def condition(self, fragment, stepList, key, hasHarvest):
+        return (fragment=="ZEE_14" or 'Eta1p7_2p7' in fragment) and '2026' in key
+upgradeWFs['ticl_v5_superclustering_mustache_ticl'] = UpgradeWorkflow_ticl_v5_superclustering(
+    steps = [
+        'RecoGlobal',
+        'HARVESTGlobal'
+    ],
+    PU = [
+        'RecoGlobal',
+        'HARVESTGlobal'
+    ],
+    suffix = '_ticl_v5_mustache',
+    offset = 0.204,
+)
+upgradeWFs['ticl_v5_superclustering_mustache_ticl'].step3 = {'--procModifiers': 'ticl_v5,ticl_superclustering_mustache_ticl'}
+upgradeWFs['ticl_v5_superclustering_mustache_ticl'].step4 = {'--procModifiers': 'ticl_v5,ticl_superclustering_mustache_ticl'}
+
+upgradeWFs['ticl_v5_superclustering_mustache_pf'] = UpgradeWorkflow_ticl_v5_superclustering(
+    steps = [
+        'RecoGlobal',
+        'HARVESTGlobal'
+    ],
+    PU = [
+        'RecoGlobal',
+        'HARVESTGlobal'
+    ],
+    suffix = '_ticl_v5_mustache_pf',
+    offset = 0.205,
+)
+upgradeWFs['ticl_v5_superclustering_mustache_pf'].step3 = {'--procModifiers': 'ticl_v5,ticl_superclustering_mustache_pf'}
+upgradeWFs['ticl_v5_superclustering_mustache_pf'].step4 = {'--procModifiers': 'ticl_v5,ticl_superclustering_mustache_pf'}
+
 # Track DNN workflows
 class UpgradeWorkflow_trackdnn(UpgradeWorkflow):
     def setup_(self, step, stepName, stepDict, k, properties):

RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py

@@ -36,6 +36,7 @@
             'keep *_ticlTrackstersCLUE3DEM_*_*',
             'keep *_ticlTrackstersCLUE3DHAD_*_*',
             'keep *_ticlTracksterLinks_*_*',
+            'keep *_ticlTracksterLinksSuperclustering_*_*',
             'keep *_ticlCandidate_*_*',
         ]
     )
@@ -60,6 +61,8 @@
       'keep *_ticlSimTracksters_*_*',
       'keep *_ticlSimTICLCandidates_*_*',
       'keep *_ticlSimTrackstersFromCP_*_*',
+      'keep *_tracksterSimTracksterAssociationLinkingSuperclustering_*_*',
+      'keep *_tracksterSimTracksterAssociationPRSuperclustering_*_*', 
       )
     )
 

RecoHGCal/TICL/interface/SuperclusteringDNNInputs.h

@@ -0,0 +1,93 @@
+/** Computation of input features for superclustering DNN. Used by plugins/TracksterLinkingBySuperClustering.cc and plugins/SuperclusteringSampleDumper.cc */
+// Author: Theo Cuisset - [email protected]
+// Date: 11/2023
+
+#ifndef __RecoHGCal_TICL_SuperclusteringDNNInputs_H__
+#define __RecoHGCal_TICL_SuperclusteringDNNInputs_H__
+
+#include <vector>
+#include <string>
+#include <memory>
+
+namespace ticl {
+  class Trackster;
+
+  // Abstract base class for DNN input preparation.
+  class AbstractSuperclusteringDNNInput {
+  public:
+    virtual ~AbstractSuperclusteringDNNInput() = default;
+
+    virtual unsigned int featureCount() const { return featureNames().size(); };
+
+    /** Get name of features. Used for SuperclusteringSampleDumper branch names (inference does not use the names, only the indices) 
+     * The default implementation is meant to be overriden by inheriting classes
+    */
+    virtual std::vector<std::string> featureNames() const {
+      std::vector<std::string> defaultNames;
+      defaultNames.reserve(featureCount());
+      for (unsigned int i = 1; i <= featureCount(); i++) {
+        defaultNames.push_back(std::string("nb_") + std::to_string(i));
+      }
+      return defaultNames;
+    }
+
+    /** Compute feature for seed and candidate pair */
+    virtual std::vector<float> computeVector(ticl::Trackster const& ts_base, ticl::Trackster const& ts_toCluster) = 0;
+  };
+
+  /* First version of DNN by Alessandro Tarabini. Meant as a DNN equivalent of Mustache algorithm (superclustering algo in ECAL)
+  Uses features : ['DeltaEta', 'DeltaPhi', 'multi_en', 'multi_eta', 'multi_pt', 'seedEta','seedPhi','seedEn', 'seedPt']
+  */
+  class SuperclusteringDNNInputV1 : public AbstractSuperclusteringDNNInput {
+  public:
+    unsigned int featureCount() const override { return 9; }
+
+    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"};
+    }
+  };
+
+  /* Second version of DNN by Alessandro Tarabini, making use of HGCAL-specific features.
+  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']
+  */
+  class SuperclusteringDNNInputV2 : public AbstractSuperclusteringDNNInput {
+  public:
+    unsigned int featureCount() const override { return 17; }
+
+    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"};
+    }
+  };
+
+  std::unique_ptr<AbstractSuperclusteringDNNInput> makeSuperclusteringDNNInputFromString(std::string dnnVersion);
+}  // namespace ticl
+
+#endif

RecoHGCal/TICL/interface/TracksterLinkingAlgoBase.h

@@ -30,11 +30,22 @@ namespace edm {
   class EventSetup;
 }  // namespace edm
 
+namespace cms {
+  namespace Ort {
+    class ONNXRuntime;
+  }
+}  // namespace cms
+
 namespace ticl {
   class TracksterLinkingAlgoBase {
   public:
-    TracksterLinkingAlgoBase(const edm::ParameterSet& conf, edm::ConsumesCollector)
-        : algo_verbosity_(conf.getParameter<int>("algo_verbosity")) {}
+    /** \param conf the configuration of the plugin
+     * \param onnxRuntime the ONNXRuntime, if onnxModelPath was provided in plugin configuration (nullptr otherwise)
+    */
+    TracksterLinkingAlgoBase(const edm::ParameterSet& conf,
+                             edm::ConsumesCollector,
+                             cms::Ort::ONNXRuntime const* onnxRuntime = nullptr)
+        : algo_verbosity_(conf.getParameter<int>("algo_verbosity")), onnxRuntime_(onnxRuntime) {}
     virtual ~TracksterLinkingAlgoBase(){};
 
     struct Inputs {
@@ -62,10 +73,14 @@ namespace ticl {
                             const edm::ESHandle<MagneticField> bfieldH,
                             const edm::ESHandle<Propagator> propH) = 0;
 
+    // To be called by TracksterLinksProducer at the start of TracksterLinksProducer::produce. Subclasses can use this to store Event and EventSetup
+    virtual void setEvent(edm::Event& iEvent, edm::EventSetup const& iEventSetup){};
+
     static void fillPSetDescription(edm::ParameterSetDescription& desc) { desc.add<int>("algo_verbosity", 0); };
 
   protected:
     int algo_verbosity_;
+    cms::Ort::ONNXRuntime const* onnxRuntime_;
   };
 }  // namespace ticl
 

RecoHGCal/TICL/plugins/BuildFile.xml

@@ -23,6 +23,8 @@
 <use name="MagneticField/Records"/>
 <use name="PhysicsTools/TensorFlow"/>
 <use name="PhysicsTools/UtilAlgos"/>
+<use name="PhysicsTools/ONNXRuntime" />
+<use name="RecoEcal/EgammaCoreTools" />
 <use name="RecoLocalCalo/HGCalRecAlgos"/>
 <use name="RecoLocalCalo/HGCalRecProducers"/>
 <use name="RecoHGCal/TICL"/>