New energy regression model

.gitignore

@@ -4,4 +4,4 @@ __init__.py
 .#*
 #*#
 *~
-*.pb
+

RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py

@@ -20,7 +20,8 @@
        'keep *_ticlTrackstersHFNoseHAD_*_*',
        'keep *_ticlTrackstersHFNoseMerge_*_*',] +
       ['keep *_pfTICL_*_*'] +
-      ['keep *_ticlGraph_*_*']
+      [ 'keep *_ticlGraph_*_*'] +
+      ['keep *_muons1stStep_*_*']
       )
     )
 TICL_RECO.outputCommands.extend(TICL_AOD.outputCommands)
@@ -48,6 +49,7 @@ def cleanOutputAndSet(outputModule, ticl_outputCommads):
                                             'keep SimVertexs_g4SimHits_*_*',
                                             'keep *_layerClusterSimClusterAssociationProducer_*_*',
                                             'keep *_layerClusterCaloParticleAssociationProducer_*_*',
+                                            'keep *_randomEngineStateProducer_*_*',
                                             ])
 
     if hasattr(process, 'FEVTDEBUGEventContent'):

RecoHGCal/TICL/interface/PatternRecognitionAlgoBase.h

@@ -38,6 +38,7 @@ namespace ticl {
       const TILES& tiles;
       const std::vector<TICLSeedingRegion>& regions;
       const tensorflow::Session* tfSession;
+      const tensorflow::Session* tfSessionER;
 
       Inputs(const edm::Event& eV,
              const edm::EventSetup& eS,
@@ -46,8 +47,9 @@ namespace ticl {
              const edm::ValueMap<std::pair<float, float>>& lT,
              const TILES& tL,
              const std::vector<TICLSeedingRegion>& rG,
-             const tensorflow::Session* tS)
-          : ev(eV), es(eS), layerClusters(lC), mask(mS), layerClustersTime(lT), tiles(tL), regions(rG), tfSession(tS) {}
+             const tensorflow::Session* tS,
+             const tensorflow::Session* tSER)
+          : ev(eV), es(eS), layerClusters(lC), mask(mS), layerClustersTime(lT), tiles(tL), regions(rG), tfSession(tS), tfSessionER(tSER) {}
     };
 
     virtual void makeTracksters(const Inputs& input,

RecoHGCal/TICL/interface/commons.h

@@ -4,6 +4,7 @@
 #include "DataFormats/Common/interface/Ref.h"
 #include "DataFormats/CaloRecHit/interface/CaloClusterFwd.h"
 #include "DataFormats/HGCalReco/interface/Trackster.h"
+#include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
 
 namespace ticl {
 
@@ -70,6 +71,64 @@ namespace ticl {
     result.emplace_back(tmpTrackster);
   }
 
+  enum LayerType {
+
+    CE_E_120 = 0,
+    CE_E_200 = 1,
+    CE_E_300 = 2,
+    CE_H_120_F = 3,
+    CE_H_200_F = 4,
+    CE_H_300_F = 5,
+    CE_H_120_C = 6,
+    CE_H_200_C = 7,
+    CE_H_300_C = 8,
+    CE_H_SCINT_C = 9,
+    EnumSize = 10
+
+  };
+
+  inline int returnIndex(DetId& lc_seed, const hgcal::RecHitTools& rhtools_) {
+    auto layer_number = rhtools_.getLayerWithOffset(lc_seed);
+    auto thickness = rhtools_.getSiThickIndex(lc_seed);
+    auto isEELayer = (layer_number <= rhtools_.lastLayerEE(false));
+    auto isScintillator = rhtools_.isScintillator(lc_seed);
+    auto isFine = (layer_number <= rhtools_.lastLayerEE(false) + 7);
+
+    if (isEELayer) {
+      if (thickness == 0) {
+        return CE_E_120;
+      } else if (thickness == 1) {
+        return CE_E_200;
+      } else if (thickness == 2) {
+        return CE_E_300;
+      }
+    } else if (!isEELayer) {
+      if (isScintillator) {
+        return CE_H_SCINT_C;
+      } else {
+        if (isFine) {
+          if (thickness == 0) {
+            return CE_H_120_F;
+          } else if (thickness == 1) {
+            return CE_H_200_F;
+          } else if (thickness == 2) {
+            return CE_H_300_F;
+          }
+        } else {
+          if (thickness == 0) {
+            return CE_H_120_C;
+          } else if (thickness == 1) {
+            return CE_H_200_C;
+          }
+            else if (thickness == 2) {
+              return CE_H_300_C;
+          }
+        }
+      }
+    }
+    return -1;
+  };
+
 }  // namespace ticl
 
 #endif

RecoHGCal/TICL/plugins/BuildFile.xml

@@ -30,6 +30,7 @@
 <use name="TrackingTools/GeomPropagators"/>
 <use name="TrackingTools/Records"/>
 <use name="TrackingTools/TrajectoryState"/>
+<use name="CommonTools/UtilAlgos"/>
 <use name="fastjet"/>
 <library file="*.cc" name="RecoHGCalTICLPlugins">
   <flags EDM_PLUGIN="1"/>

RecoHGCal/TICL/plugins/PatternRecognitionbyCLUE3D.cc

@@ -9,6 +9,7 @@
 #include "FWCore/MessageLogger/interface/MessageLogger.h"
 #include "FWCore/Utilities/interface/Exception.h"
 #include "PatternRecognitionbyCLUE3D.h"
+#include "RecoHGCal/TICL/interface/commons.h"
 
 #include "TrackstersPCA.h"
 #include "Geometry/CaloGeometry/interface/CaloGeometry.h"
@@ -38,7 +39,9 @@ PatternRecognitionbyCLUE3D<TILES>::PatternRecognitionbyCLUE3D(const edm::Paramet
       outlierMultiplier_(conf.getParameter<double>("outlierMultiplier")),
       minNumLayerCluster_(conf.getParameter<int>("minNumLayerCluster")),
       eidInputName_(conf.getParameter<std::string>("eid_input_name")),
+      eidInputNameER_(conf.getParameter<std::string>("eid_input_nameER")),
       eidOutputNameEnergy_(conf.getParameter<std::string>("eid_output_name_energy")),
+      eidOutputNameEnergyER_(conf.getParameter<std::string>("eid_output_name_energyER")),
       eidOutputNameId_(conf.getParameter<std::string>("eid_output_name_id")),
       eidMinClusterEnergy_(conf.getParameter<double>("eid_min_cluster_energy")),
       eidNLayers_(conf.getParameter<int>("eid_n_layers")),
@@ -327,7 +330,7 @@ void PatternRecognitionbyCLUE3D<TILES>::makeTracksters(
                               rhtools_.getPositionLayer(rhtools_.lastLayerEE(false), false).z());
 
   // run energy regression and ID
-  energyRegressionAndID(input.layerClusters, input.tfSession, result);
+  energyRegressionAndID(input.layerClusters, input.tfSession, input.tfSessionER, result);
   if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > PatternRecognitionAlgoBaseT<TILES>::Advanced) {
     for (auto const &t : result) {
       edm::LogVerbatim("PatternRecognitionbyCLUE3D") << "Barycenter: " << t.barycenter();
@@ -352,6 +355,7 @@ void PatternRecognitionbyCLUE3D<TILES>::makeTracksters(
 template <typename TILES>
 void PatternRecognitionbyCLUE3D<TILES>::energyRegressionAndID(const std::vector<reco::CaloCluster> &layerClusters,
                                                               const tensorflow::Session *eidSession,
+                                                              const tensorflow::Session *eidSessionEnergyRegression,
                                                               std::vector<Trackster> &tracksters) {
   // Energy regression and particle identification strategy:
   //
@@ -418,7 +422,6 @@ void PatternRecognitionbyCLUE3D<TILES>::energyRegressionAndID(const std::vector<
   // fill input tensor (5)
   for (int i = 0; i < batchSize; i++) {
     const Trackster &trackster = tracksters[tracksterIndices[i]];
-
     // per layer, we only consider the first eidNClusters_ clusters in terms of energy, so in order
     // to avoid creating large / nested structures to do the sorting for an unknown number of total
     // clusters, create a sorted list of layer cluster indices to keep track of the filled clusters
@@ -466,16 +469,6 @@ void PatternRecognitionbyCLUE3D<TILES>::energyRegressionAndID(const std::vector<
   // run the inference (7)
   tensorflow::run(const_cast<tensorflow::Session *>(eidSession), inputList, outputNames, &outputs);
 
-  // store regressed energy per trackster (8)
-  if (!eidOutputNameEnergy_.empty()) {
-    // get the pointer to the energy tensor, dimension is batch x 1
-    float *energy = outputs[0].flat<float>().data();
-
-    for (const int &i : tracksterIndices) {
-      tracksters[i].setRegressedEnergy(*(energy++));
-    }
-  }
-
   // store id probabilities per trackster (8)
   if (!eidOutputNameId_.empty()) {
     // get the pointer to the id probability tensor, dimension is batch x id_probabilities.size()
@@ -487,6 +480,69 @@ void PatternRecognitionbyCLUE3D<TILES>::energyRegressionAndID(const std::vector<
       probs += tracksters[i].id_probabilities().size();
     }
   }
+  // NEW ENERGY REGRESSION!
+  std::vector<int> tracksterIndicesHadronic;
+  auto index_trackster = 0;
+  for (auto &t : tracksters) {
+    t.setRegressedEnergy(t.raw_energy());
+    if ((t.id_probability(ticl::Trackster::ParticleType::photon) +
+         t.id_probability(ticl::Trackster::ParticleType::electron)) >= 0.) {
+      tracksterIndicesHadronic.push_back(index_trackster);
+    }
+    index_trackster++;
+  }
+  batchSize = static_cast<int>(tracksterIndicesHadronic.size());
+  if (batchSize == 0) {
+    return;
+  }
+  // create input and output tensors for energy regression
+  tensorflow::TensorShape shapeER({batchSize, eidNFeaturesER_});
+  tensorflow::Tensor inputEnergyRegression(tensorflow::DT_FLOAT, shapeER);
+  tensorflow::NamedTensorList inputListEnergyRegression = {{eidInputNameER_, inputEnergyRegression}};
+
+  std::vector<tensorflow::Tensor> outputsEnergyRegression;
+  std::vector<std::string> outputNamesEnergyRegression;
+
+  outputNamesEnergyRegression.push_back(eidOutputNameEnergyER_);
+
+  // fill input tensor (5)
+  for (int i = 0; i < batchSize; i++) {
+    const Trackster &trackster = tracksters[tracksterIndicesHadronic[i]]; 
+    float en_total = 0;
+    float *featuresER = &inputEnergyRegression.tensor<float, 2>()(i, 0);
+    std::array<float, eidNFeaturesER_> inputsArrayER;
+    for (size_t j = 0; j < inputsArrayER.size(); j++) {
+      inputsArrayER[j] = 0.;
+    }
+
+    for (int k = 0; k < (int)trackster.vertices().size(); k++) {
+      const reco::CaloCluster &cluster = layerClusters[trackster.vertices(k)];
+      auto lc_seed = cluster.seed();
+      auto index = returnIndex(lc_seed, rhtools_);
+        if(index >= 0){
+          inputsArrayER[index] += cluster.energy() / trackster.raw_energy(); //energy fraction
+        }
+      en_total += cluster.energy();
+    }
+
+    inputsArrayER[ticl::LayerType::EnumSize] = trackster.barycenter().eta(); //eta
+    inputsArrayER[ticl::LayerType::EnumSize+1] = trackster.raw_energy(); //total trackster energy
+    for(size_t i_p = 0; i_p < trackster.id_probabilities().size(); i_p++){
+      inputsArrayER[ticl::LayerType::EnumSize+2+i_p] = trackster.id_probabilities(i_p); //probabilities
+    }
+    for (size_t j = 0; j < inputsArrayER.size(); j++) {
+      *(featuresER++) = inputsArrayER[j];
+    }
+  }
+
+  tensorflow::run(const_cast<tensorflow::Session *>(eidSessionEnergyRegression),
+                  inputListEnergyRegression,
+                  outputNamesEnergyRegression,
+                  &outputsEnergyRegression);
+  float *energy = outputsEnergyRegression[0].flat<float>().data();
+  for (const int &i : tracksterIndicesHadronic) {
+    tracksters[i].setRegressedEnergy(*(energy++));
+  }
 }
 
 template <typename TILES>
@@ -869,7 +925,9 @@ void PatternRecognitionbyCLUE3D<TILES>::fillPSetDescription(edm::ParameterSetDes
   iDesc.add<double>("outlierMultiplier", 2);
   iDesc.add<int>("minNumLayerCluster", 2)->setComment("Not Inclusive");
   iDesc.add<std::string>("eid_input_name", "input");
+  iDesc.add<std::string>("eid_input_nameER", "input");
   iDesc.add<std::string>("eid_output_name_energy", "output/regressed_energy");
+  iDesc.add<std::string>("eid_output_name_energyER", "output/er_coefficients");
   iDesc.add<std::string>("eid_output_name_id", "output/id_probabilities");
   iDesc.add<double>("eid_min_cluster_energy", 1.);
   iDesc.add<int>("eid_n_layers", 50);

RecoHGCal/TICL/plugins/PatternRecognitionbyCLUE3D.h

@@ -5,6 +5,7 @@
 #define __RecoHGCal_TICL_PRbyCLUE3D_H__
 #include <memory>  // unique_ptr
 #include "RecoHGCal/TICL/interface/PatternRecognitionAlgoBase.h"
+#include "DataFormats/DetId/interface/DetId.h"
 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
 
 namespace ticl {
@@ -19,6 +20,7 @@ namespace ticl {
                         std::unordered_map<int, std::vector<int>>& seedToTracksterAssociation) override;
 
     void energyRegressionAndID(const std::vector<reco::CaloCluster>& layerClusters,
+                               const tensorflow::Session*,
                                const tensorflow::Session*,
                                std::vector<Trackster>& result);
 
@@ -129,16 +131,20 @@ namespace ticl {
     const int minNumLayerCluster_;
     const std::vector<int> filter_on_categories_;
     const std::string eidInputName_;
+    const std::string eidInputNameER_;
     const std::string eidOutputNameEnergy_;
+    const std::string eidOutputNameEnergyER_;
     const std::string eidOutputNameId_;
     const float eidMinClusterEnergy_;
     const int eidNLayers_;
     const int eidNClusters_;
 
     hgcal::RecHitTools rhtools_;
     tensorflow::Session* eidSession_;
+    tensorflow::Session* eidSessionER_;
 
     static const int eidNFeatures_ = 3;
+    static const int eidNFeaturesER_ = 20;
   };
 
 }  // namespace ticl