Merge branch 'pflow_NN' into '23.0'

Add optional ML based e/p workflow

See merge request !62022

Reconstruction/eflowRec/CMakeLists.txt

@@ -5,6 +5,7 @@ atlas_subdir( eflowRec )
 
 # External dependencies:
 find_package( ROOT COMPONENTS Core Tree MathCore Hist RIO pthread )
+find_package( onnxruntime )
 
 # Component(s) in the package:
 atlas_add_component( eflowRec
@@ -14,7 +15,7 @@ atlas_add_component( eflowRec
                      LINK_LIBRARIES ${ROOT_LIBRARIES} CaloEvent CaloIdentifier AthContainers AthLinks AthenaBaseComps CxxUtils AthenaKernel GeoPrimitives Identifier xAODBase 
 		     xAODCaloEvent xAODCore xAODEgamma xAODMuon xAODPFlow xAODTau xAODTracking GaudiKernel InDetReadoutGeometry TRT_ReadoutGeometry Particle RecoToolInterfaces 
 		     TrkParameters CaloDetDescrLib CaloUtilsLib StoreGateLib FourMomUtils PathResolver TrkCaloExtension TrkParametersIdentificationHelpers FourMomUtils 
-		     InDetTrackSelectionToolLib AthenaMonitoringKernelLib ICaloTrkMuIdTools AsgMessagingLib)
+		     InDetTrackSelectionToolLib AthenaMonitoringKernelLib ICaloTrkMuIdTools AsgMessagingLib AthOnnxruntimeServiceLib)
 
 # Install files from the package:
 atlas_install_python_modules( python/*.py POST_BUILD_CMD ${ATLAS_FLAKE8} )

Reconstruction/eflowRec/eflowRec/PFEnergyPredictorTool.h

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef PFENERFYPREDICTORTOOL_H
+#define PFENERFYPREDICTORTOOL_H
+
+#include "AthenaBaseComps/AthAlgTool.h"
+#include "GaudiKernel/ServiceHandle.h"
+#include "AthOnnxruntimeService/IONNXRuntimeSvc.h"
+#include <fstream>      // std::fstream
+
+static const InterfaceID IID_PFEnergyPredictorTool("PFEnergyPredictorTool", 1, 0);
+class eflowRecTrack;
+
+class PFEnergyPredictorTool :  public AthAlgTool
+{
+public:
+  PFEnergyPredictorTool(const std::string& type, const std::string& name, const IInterface* parent);
+  virtual StatusCode initialize() override;
+  virtual StatusCode finalize()   override;
+
+  float runOnnxInference(std::vector<float> &tensor) const;
+  static const InterfaceID& interfaceID();
+
+  float nnEnergyPrediction(const eflowRecTrack *ptr) const;
+  void NormalizeTensor(std::vector<float> &tensor, size_t limit) const;
+
+private:
+  //mark as thread safe because we need to call the run function of Session, which is not const
+  //the onnx documentation states that this is thread safe
+  std::unique_ptr<Ort::Session> m_session ATLAS_THREAD_SAFE;
+
+  std::vector<const char *> m_input_node_names;
+
+  std::vector<const char *> m_output_node_names;
+
+  std::vector<int64_t> m_input_node_dims;
+  ServiceHandle<AthONNX::IONNXRuntimeSvc> m_svc{this, "ONNXRuntimeSvc", "AthONNX::ONNXRuntimeSvc", "CaloMuonScoreTool ONNXRuntimeSvc"};
+  Gaudi::Property<std::string> m_model_filepath{this, "ModelPath", "////"};
+
+  /** Normalization constants for the inputs to the onnx model */
+  Gaudi::Property<float> m_cellE_mean{this,"cellE_mean",-2.2852574689444385};
+  Gaudi::Property<float> m_cellE_std{this,"cellE_std",2.0100506557174946};
+  Gaudi::Property<float> m_cellPhi_std{this,"cellPhi_std",0.6916977411859621};
+  
+};
+
+inline const InterfaceID& PFEnergyPredictorTool::interfaceID() { return IID_PFEnergyPredictorTool; }
+
+
+#endif
+

Reconstruction/eflowRec/eflowRec/PFSubtractionTool.h

@@ -12,6 +12,7 @@
 #include "eflowRec/EtaPhiLUT.h"
 #include "eflowRec/IPFSubtractionTool.h"
 #include "eflowRec/PFData.h"
+#include "eflowRec/PFEnergyPredictorTool.h"
 #include "eflowRec/PFMatchPositions.h"
 #include "eflowRec/PFTrackClusterMatchingTool.h"
 #include "eflowRec/PFCalcRadialEnergyProfiles.h"
@@ -94,6 +95,14 @@ private:
   PFSubtractionStatusSetter m_pfSubtractionStatusSetter{};
   PFSubtractionEnergyRatioCalculator m_pfSubtractionEnergyRatioCalculator{};
   eflowSubtract::Subtractor m_subtractor{};
+
+  /** Tool for getting predictiing the energy using an ONNX model */
+  ToolHandle<PFEnergyPredictorTool> m_NNEnergyPredictorTool{this, "NNEnergyPredictorTool", "","Tool for getting predictiing the energy using an ONNX model "};
+  
+  /** Toggle whether we use the neural net energy */
+  Gaudi::Property<bool> m_useNNEnergy{this, "useNNEnergy", false, "Toggle whether we use the neural net energy"};
+
+
 };
 
 #endif

Reconstruction/eflowRec/eflowRec/eflowCaloObject.h

@@ -25,6 +25,7 @@ class eflowRecTrack;
 class eflowTrackClusterLink;
 class eflowLayerIntegrator;
 class eflowEEtaBinnedParameters;
+class PFEnergyPredictorTool;
 
 /**
 An internal EDM object which stores information about systems of associated tracks and calorimeter clusters. Specifically it stores vectors of pointers to eflowRecTracks, eflowRecClusters and eflowTrackClusterLinks. In addition it stores links to an xAOD::CaloClusterContainer and its associated aux container. This class also calculates the expected energy deposit in the calorimeter from a track in the system, and stores that information so that clients can retrieve it. It also calculates the calorimeter cell ordering to be used in the subtraction. Both of these things are done in the simulateShower method which uses the data stored in an eflowEEtaBinnedParameters object, which is filled by e.g the eflowCellEOverP_mc12_JetETMiss tool.
@@ -68,7 +69,8 @@ public:
   double getExpectedVariance() const;
   double getClusterEnergy() const ;
 
-  void simulateShower(eflowLayerIntegrator *integrator, const eflowEEtaBinnedParameters* binnedParameters, bool useUpdated2015ChargedShowerSubtraction);
+  void simulateShower(eflowLayerIntegrator *integrator, const eflowEEtaBinnedParameters* binnedParameters, bool useUpdated2015ChargedShowerSubtraction,
+  const PFEnergyPredictorTool* energyP);
 
 private:
 

Reconstruction/eflowRec/eflowRec/eflowTrackCaloPoints.h

@@ -31,15 +31,21 @@ class eflowTrackCaloPoints {
  public:
 
   eflowTrackCaloPoints(const std::map<eflowCalo::LAYER, const Trk::TrackParameters*> & trackParameters);
+  eflowTrackCaloPoints(const std::map<eflowCalo::LAYER, const Trk::TrackParameters*> & trackParameters,
+  std::map<CaloCell_ID::CaloSample,const Trk::TrackParameters*>& tileTrackParamaters);
  eflowTrackCaloPoints() : m_isEM1Barrel(false), m_isEM2Barrel(false)  {}
   ~eflowTrackCaloPoints();
 
   const std::pair<float, float> operator[] (eflowCalo::LAYER layer) const;
   const eflowEtaPhiPosition& getEtaPhiPos(eflowCalo::LAYER layer) const;
+  const eflowEtaPhiPosition& getTileEtaPhiPos(CaloCell_ID::CaloSample layer) const;
 
   double getEta(eflowCalo::LAYER layer) const {return getEtaPhiPos(layer).getEta();}
   double getPhi(eflowCalo::LAYER layer) const {return getEtaPhiPos(layer).getPhiD();}
 
+  double getTileEta(CaloCell_ID::CaloSample layer) const {return getTileEtaPhiPos(layer).getEta();}
+  double getTilePhi(CaloCell_ID::CaloSample layer) const {return getTileEtaPhiPos(layer).getPhiD();}
+
   const std::pair<float, float> getEM2etaPhi() const  {return (*this)[getEM2Layer()]; }
   const eflowEtaPhiPosition& getEM2etaPhiPos() const  {return getEtaPhiPos(getEM2Layer()); }
   double getEM2eta() const {return getEM2etaPhiPos().getEta(); }
@@ -55,6 +61,7 @@ class eflowTrackCaloPoints {
   inline bool haveLayer(eflowCalo::LAYER layer) const { return getEta(layer) != m_defaultEtaPhiPair.first;  }
 
   void setEtaPhi(eflowCaloENUM secondLayer, double eta, double phi);
+  void setEtaPhiTile(CaloCell_ID::CaloSample secondLayer, const Amg::Vector3D& vec);
   void setEtaPhi(eflowCalo::LAYER lay, const Amg::Vector3D& vec);
   void copyEtaPhi(eflowCalo::LAYER fromLay, eflowCalo::LAYER toLay);
 
@@ -78,6 +85,7 @@ class eflowTrackCaloPoints {
   std::map<eflowCalo::LAYER, Amg::Vector3D > m_positions;
   std::map<eflowCalo::LAYER, Amg::Vector3D > m_directions;
   std::map<eflowCalo::LAYER, eflowEtaPhiPosition>  m_etaPhiPositions;
+  std::map<CaloCell_ID::CaloSample,eflowEtaPhiPosition> m_tileEtaPhiPositions;
 };
 
 inline void eflowTrackCaloPoints::copyEtaPhi(eflowCalo::LAYER fromLay, eflowCalo::LAYER toLay) {

Reconstruction/eflowRec/python/PFCfg.py

@@ -61,7 +61,7 @@ def getPFTrackClusterMatchingTool(inputFlags,matchCut,distanceType,clusterPositi
 
 def getPFCellLevelSubtractionTool(inputFlags,toolName):
     PFCellLevelSubtractionToolFactory = CompFactory.PFSubtractionTool
-    PFCellLevelSubtractionTool = PFCellLevelSubtractionToolFactory(toolName)
+    PFCellLevelSubtractionTool = PFCellLevelSubtractionToolFactory(toolName,useNNEnergy = inputFlags.PF.useMLEOverP)
 
     if inputFlags.GeoModel.Run <= LHCPeriod.Run3:
         eflowCellEOverPTool_Run2_mc20_JetETMiss = CompFactory.eflowCellEOverPTool_Run2_mc20_JetETMiss
@@ -84,11 +84,15 @@ def getPFCellLevelSubtractionTool(inputFlags,toolName):
     PFCellLevelSubtractionTool.PFTrackClusterMatchingTool_015 = getPFTrackClusterMatchingTool(inputFlags,0.15,"EtaPhiSquareDistance","PlainEtaPhi","MatchingTool_Pull_015")
     PFCellLevelSubtractionTool.PFTrackClusterMatchingTool_02 = getPFTrackClusterMatchingTool(inputFlags,0.2,"EtaPhiSquareDistance","PlainEtaPhi","MatchingTool_Pull_02")
 
+    if inputFlags.PF.useMLEOverP:
+        PFEnergyPredictorTool = CompFactory.PFEnergyPredictorTool("PFCellLevelEnergyPredcictorTool",ModelPath = inputFlags.PF.EOverP_NN_Model)
+        PFCellLevelSubtractionTool.NNEnergyPredictorTool = PFEnergyPredictorTool
+
     return PFCellLevelSubtractionTool
 
 def getPFRecoverSplitShowersTool(inputFlags,toolName):
     PFRecoverSplitShowersToolFactory = CompFactory.PFSubtractionTool
-    PFRecoverSplitShowersTool = PFRecoverSplitShowersToolFactory(toolName)
+    PFRecoverSplitShowersTool = PFRecoverSplitShowersToolFactory(toolName,useNNEnergy = inputFlags.PF.useMLEOverP)
 
     if inputFlags.GeoModel.Run <= LHCPeriod.Run3:
         eflowCellEOverPTool_Run2_mc20_JetETMiss = CompFactory.eflowCellEOverPTool_Run2_mc20_JetETMiss
@@ -99,6 +103,10 @@ def getPFRecoverSplitShowersTool(inputFlags,toolName):
 
     PFRecoverSplitShowersTool.RecoverSplitShowers = True
 
+    if inputFlags.PF.useMLEOverP:
+        PFEnergyPredictorTool = CompFactory.PFEnergyPredictorTool("PFRecoverSplitShowersEnergyPredcictorTool",ModelPath = inputFlags.PF.EOverP_NN_Model)
+        PFRecoverSplitShowersTool.NNEnergyPredictorTool = PFEnergyPredictorTool
+
     return PFRecoverSplitShowersTool
 
 def getPFMomentCalculatorTool(inputFlags, momentsToCalculateList):

Reconstruction/eflowRec/python/PFConfigFlags.py

@@ -15,5 +15,7 @@ def createPFConfigFlags():
     pfConfigFlags.addFlag("PF.useMuLinks", lambda prevFlags : prevFlags.Reco.EnableCombinedMuon  or prevFlags.PF.useRecExCommon)
     pfConfigFlags.addFlag("PF.useOldPFO",False)
     pfConfigFlags.addFlag("PF.useRecExCommon",False) #Toggle whether we are in the RecExCommon config or not.
+    pfConfigFlags.addFlag("PF.useMLEOverP",False) #Toggle whether to use the Machine Learning based EOverP inference or not
+    pfConfigFlags.addFlag("PF.EOverP_NN_Model",'/afs/cern.ch/user/m/mhodgkin/onnx_15_03_23.onnx') #Model to use in EOverP inference
     
     return pfConfigFlags

Reconstruction/eflowRec/src/PFEnergyPredictorTool.cxx

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#include "eflowRec/PFEnergyPredictorTool.h"
+#include "eflowRec/eflowCaloObject.h"
+#include "eflowRec/eflowTrackClusterLink.h"
+#include "eflowRec/eflowRecCluster.h"
+
+#include "CaloGeoHelpers/CaloSampling.h"
+
+PFEnergyPredictorTool::PFEnergyPredictorTool(const std::string& type, const std::string& name, const IInterface* parent) : AthAlgTool(type, name, parent)
+{
+
+}
+
+
+StatusCode PFEnergyPredictorTool::initialize()
+{
+   ATH_MSG_DEBUG("Initializing " << name());
+   if(m_model_filepath == "////"){
+      ATH_MSG_WARNING("model not provided tool will not work");
+      return StatusCode::SUCCESS;
+   }
+   ATH_CHECK(m_svc.retrieve());
+   std::string path  = m_model_filepath;//Add path resolving code
+
+   Ort::SessionOptions session_options;
+   Ort::AllocatorWithDefaultOptions allocator;
+   session_options.SetIntraOpNumThreads(1);
+   session_options.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
+   m_session = std::make_unique<Ort::Session>(m_svc->env(), path.c_str(), session_options);
+
+    ATH_MSG_INFO("Created ONNX runtime session with model " << path);
+
+    size_t num_input_nodes = m_session->GetInputCount();
+    m_input_node_names.resize(num_input_nodes);
+
+    for (std::size_t i = 0; i < num_input_nodes; i++) {
+        // print input node names
+        char *input_name = m_session->GetInputName(i, allocator);
+        ATH_MSG_INFO("Input " << i << " : "
+                              << " name= " << input_name);
+        m_input_node_names[i] = input_name;
+        // print input node types
+        Ort::TypeInfo type_info = m_session->GetInputTypeInfo(i);
+        auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
+        ONNXTensorElementDataType type = tensor_info.GetElementType();
+        ATH_MSG_INFO("Input " << i << " : "
+                              << " type= " << type);
+
+        // print input shapes/dims
+        m_input_node_dims = tensor_info.GetShape();
+        m_input_node_dims[1] = 5430/5;
+        ATH_MSG_INFO("Input " << i << " : num_dims= " << m_input_node_dims.size());
+        for (std::size_t j = 0; j < m_input_node_dims.size(); j++) {
+            if (m_input_node_dims[j] < 0) m_input_node_dims[j] = 1;
+            ATH_MSG_INFO("Input " << i << " : dim " << j << "= " << m_input_node_dims[j]);
+        }
+    }
+    
+    // output nodes
+    std::vector<int64_t> output_node_dims;
+    size_t num_output_nodes = m_session->GetOutputCount();
+    ATH_MSG_INFO("Have output nodes " << num_output_nodes);
+    m_output_node_names.resize(num_output_nodes);
+
+    for (std::size_t i = 0; i < num_output_nodes; i++) {
+        // print output node names
+        char *output_name = m_session->GetOutputName(i, allocator);
+        ATH_MSG_INFO("Output " << i << " : "
+                               << " name= " << output_name);
+        m_output_node_names[i] = output_name;
+
+        Ort::TypeInfo type_info = m_session->GetOutputTypeInfo(i);
+        auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
+        ONNXTensorElementDataType type = tensor_info.GetElementType();
+        ATH_MSG_INFO("Output " << i << " : "
+                               << " type= " << type);
+
+        // print output shapes/dims
+        output_node_dims = tensor_info.GetShape();
+        ATH_MSG_INFO("Output " << i << " : num_dims= " << output_node_dims.size());
+        for (std::size_t j = 0; j < output_node_dims.size(); j++) {
+            if (output_node_dims[j] < 0) output_node_dims[j] = 1;
+            ATH_MSG_INFO("Output" << i << " : dim " << j << "= " << output_node_dims[j]);
+        }
+    }
+
+    return StatusCode::SUCCESS;
+}
+
+float PFEnergyPredictorTool::runOnnxInference(std::vector<float> &tensor) const {
+    using std::endl;
+    using std::cout;
+    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
+    auto input_tensor_size = tensor.size();
+
+    Ort::Value input_tensor =
+        Ort::Value::CreateTensor<float>(memory_info, tensor.data(), input_tensor_size,
+         m_input_node_dims.data(), m_input_node_dims.size());
+
+    auto output_tensors = m_session->Run(Ort::RunOptions{nullptr}, m_input_node_names.data(), &input_tensor, m_input_node_names.size(),
+                                         m_output_node_names.data(), m_output_node_names.size());
+
+    const float *output_score_array = output_tensors.front().GetTensorData<float>();
+
+    // Binary classification - the score is just the first element of the output tensor
+    float output_score = output_score_array[0];
+
+    return output_score;
+}
+
+std::array<double,19> getEtaTrackCalo(const eflowTrackCaloPoints& trackCaloPoints) {
+    return std::array<double,19> { trackCaloPoints.getEta(eflowCalo::EMB1), trackCaloPoints.getEta(eflowCalo::EMB2), trackCaloPoints.getEta(eflowCalo::EMB3),
+    trackCaloPoints.getEta(eflowCalo::EME1), trackCaloPoints.getEta(eflowCalo::EME2), trackCaloPoints.getEta(eflowCalo::EME3),
+    trackCaloPoints.getEta(eflowCalo::HEC1), trackCaloPoints.getEta(eflowCalo::HEC2), trackCaloPoints.getEta(eflowCalo::HEC3),trackCaloPoints.getEta(eflowCalo::HEC4),    
+    trackCaloPoints.getTileEta(CaloSampling::TileBar0),trackCaloPoints.getTileEta(CaloSampling::TileBar1),trackCaloPoints.getTileEta(CaloSampling::TileBar2),
+    trackCaloPoints.getTileEta(CaloSampling::TileGap1),trackCaloPoints.getTileEta(CaloSampling::TileGap2),trackCaloPoints.getTileEta(CaloSampling::TileGap3),
+    trackCaloPoints.getTileEta(CaloSampling::TileExt0),trackCaloPoints.getTileEta(CaloSampling::TileExt1),trackCaloPoints.getTileEta(CaloSampling::TileExt2)};
+}
+
+
+std::array<double,19> getPhiTrackCalo(const eflowTrackCaloPoints& trackCaloPoints) {
+    return std::array<double,19> { trackCaloPoints.getPhi(eflowCalo::EMB1), trackCaloPoints.getPhi(eflowCalo::EMB2), trackCaloPoints.getPhi(eflowCalo::EMB3),
+    trackCaloPoints.getPhi(eflowCalo::EME1), trackCaloPoints.getPhi(eflowCalo::EME2), trackCaloPoints.getPhi(eflowCalo::EME3),
+    trackCaloPoints.getPhi(eflowCalo::HEC1), trackCaloPoints.getPhi(eflowCalo::HEC2), trackCaloPoints.getPhi(eflowCalo::HEC3),trackCaloPoints.getPhi(eflowCalo::HEC4),
+    trackCaloPoints.getTilePhi(CaloSampling::TileBar0),trackCaloPoints.getTilePhi(CaloSampling::TileBar1),trackCaloPoints.getTilePhi(CaloSampling::TileBar2),
+    trackCaloPoints.getTilePhi(CaloSampling::TileGap1),trackCaloPoints.getTilePhi(CaloSampling::TileGap2),trackCaloPoints.getTilePhi(CaloSampling::TileGap3),
+    trackCaloPoints.getTilePhi(CaloSampling::TileExt0),trackCaloPoints.getTilePhi(CaloSampling::TileExt1),trackCaloPoints.getTilePhi(CaloSampling::TileExt2)};
+}
+
+
+float PFEnergyPredictorTool::nnEnergyPrediction(const eflowRecTrack *ptr) const{
+
+     constexpr std::array<int,19> calo_numbers{1,2,3,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20};
+     constexpr std::array<int,12> fixed_r_numbers = {1,2,3,12,13,14,15,16,17,18,19,20};
+     constexpr std::array<double,12> fixed_r_vals = {1532.18, 1723.89, 1923.02, 2450.00, 2995.00, 3630.00, 3215.00,
+                                        3630.00, 2246.50, 2450.00, 2870.00, 3480.00
+                                       };
+     constexpr std::array<int, 7> fixed_z_numbers = {5,6,7,8,9,10,11};
+     constexpr std::array<double, 7> fixed_z_vals = {3790.03, 3983.68, 4195.84, 4461.25, 4869.50, 5424.50, 5905.00};
+     std::unordered_map<int, double> r_calo_dict;//change to flatmap in c++23
+     std::unordered_map<int, double> z_calo_dict;
+     for(size_t i=0; i<fixed_r_vals.size(); i++) r_calo_dict[fixed_r_numbers[i]] = fixed_r_vals[i];
+     for(size_t i=0; i<fixed_z_numbers.size(); i++) z_calo_dict[fixed_z_numbers[i]] = fixed_z_vals[i];
+
+     std::vector<float> inputnn;
+     inputnn.assign(5430, 0.0);
+     std::vector<eflowRecCluster*> matchedClusters;
+     std::vector<eflowTrackClusterLink*> links = ptr->getClusterMatches();
+
+    std::array<double, 19> etatotal = getEtaTrackCalo(ptr->getTrackCaloPoints());
+    std::array<double, 19> phitotal = getPhiTrackCalo(ptr->getTrackCaloPoints());
+
+     const std::array<double, 2> track{ptr->getTrack()->eta(), ptr->getTrack()->phi()};
+     double totalE = 0.0;
+
+     for(auto clink : links){
+        auto cell = clink->getCluster()->getCluster();
+        float clusterE = cell->e()*1e-3;
+        float clusterEta = cell->eta();
+
+        if (clusterE < 0.0 || clusterE > 1e4f || std::abs(clusterEta) > 2.5) continue;
+
+        constexpr bool cutOnR = false;
+        if(cutOnR){
+            std::array<double, 2> p{clink->getCluster()->getCluster()->eta(), clink->getCluster()->getCluster()->phi()};
+            double part1 = p[0] - track[0];
+            double part2 = p[1] - track[1];
+            while(part1 > M_PI) part1 -= 2*M_PI;
+            while(part1 < -M_PI) part1 += 2*M_PI;
+            while(part2 > M_PI) part2 -= 2*M_PI;
+            while(part2 < -M_PI) part2 += 2*M_PI;
+            double R = std::sqrt(part1 * part1 + part2*part2);
+            if(R >= 1.2) continue;
+        }
+        totalE += clusterE;
+
+        matchedClusters.push_back(clink->getCluster());
+     }
+
+     std::vector<std::array<double, 5>> cells;
+
+    const eflowTrackCaloPoints& trackCaloPoints = ptr->getTrackCaloPoints();
+    bool trk_bool_em[2] = {false,false};
+    std::array<double,2> trk_em_eta = {trackCaloPoints.getEta(eflowCalo::EMB2), trackCaloPoints.getEta(eflowCalo::EME2)};
+    std::array<double,2> trk_em_phi = {trackCaloPoints.getPhi(eflowCalo::EMB2), trackCaloPoints.getPhi(eflowCalo::EME2)};
+    double eta_ctr;
+    double phi_ctr;
+    for(int i =0; i<2; i++) {
+        trk_bool_em[i] = std::abs(trk_em_eta[i]) < 2.5 && std::abs(trk_em_phi[i]) <= M_PI;
+    }
+    int nProj_em = (int)trk_bool_em[0] + (int)trk_bool_em[1];
+
+    if(nProj_em ==1) {
+        eta_ctr = trk_bool_em[0] ? trk_em_eta[0] : trk_em_eta[1];
+        phi_ctr = trk_bool_em[0] ? trk_em_phi[0] : trk_em_phi[1];
+    } else if(nProj_em==2) {
+        eta_ctr = (trk_em_eta[0] + trk_em_eta[1]) / 2.0;
+        phi_ctr = (trk_em_phi[0] + trk_em_phi[1]) / 2.0;
+    } else {
+        eta_ctr = ptr->getTrack()->eta();
+        phi_ctr = ptr->getTrack()->phi();
+    }
+
+
+
+     for(auto cptr : matchedClusters){
+        auto clustlink = cptr->getCluster();
+
+        for(auto it_cell = clustlink->cell_begin(); it_cell != clustlink->cell_end(); it_cell++){
+           const CaloCell* cell = (*it_cell);
+           float cellE = cell->e()*(it_cell.weight())*1e-3f;
+           if(cellE < 0.005) continue;//Cut from ntuple maker
+           auto theDDE=it_cell->caloDDE();
+           double cx=theDDE->x();
+           double cy=theDDE->y();
+
+           cells.emplace_back( std::array<double, 5> { cellE,
+                    theDDE->eta() -  eta_ctr,
+                    theDDE->phi() -  phi_ctr,
+                    std::hypot(cx,cy), //rperp
+                    0.0    } );
+        }
+     }
+
+
+    std::vector<bool> trk_bool(calo_numbers.size(), false);
+    std::vector<std::array<double,4>> trk_full(calo_numbers.size());
+    for(size_t j=0; j<phitotal.size(); j++) {
+        int cnum = calo_numbers[j];
+        double eta = etatotal[j];
+        double phi = phitotal[j];
+        if(std::abs(eta) < 2.5 && std::abs(phi) <= M_PI) {
+            trk_bool[j] = true;
+            trk_full[j][0] = eta;
+            trk_full[j][1] = phi;
+            trk_full[j][3] = cnum;
+            double rPerp =-99999;
+            if(auto itr = r_calo_dict.find(cnum); itr != r_calo_dict.end()) rPerp = itr->second;
+            else if(auto itr = z_calo_dict.find(cnum); itr != z_calo_dict.end())
+            {
+                double z = itr->second;
+                if(eta != 0.0){
+                   double aeta = std::abs(eta);
+                   rPerp = z*2.*std::exp(aeta)/(std::exp(2.0*aeta)-1.0);
+                }else rPerp =0.0; //Check if this makes sense
+            } else {
+                throw std::runtime_error("Calo sample num not found in dicts..");
+            }
+            trk_full[j][2] = rPerp;
+        } else {
+            trk_full[j].fill(0.0);
+        }
+    }
+    double trackP = std::abs(1. / ptr->getTrack()->qOverP()) * 1e-3;
+    int trk_proj_num = std::accumulate(trk_bool.begin(), trk_bool.end(), 0);
+    if(trk_proj_num ==0) {
+        trk_proj_num =1;
+        std::array<double,5> trk_arr;
+
+        trk_arr[0] = trackP;
+        trk_arr[1] = ptr->getTrack()->eta() - eta_ctr;
+        trk_arr[2] = ptr->getTrack()->phi() - phi_ctr;
+        trk_arr[3] = 1532.18; // just place it in EMB1
+        trk_arr[4] = 1.;
+
+        cells.emplace_back(trk_arr);
+    } else {
+        for(size_t i =0; i<calo_numbers.size(); i++) {
+            if(trk_bool[i]==false) continue;
+            std::array<double,5> trk_arr;
+            trk_arr[0]= trackP/double(trk_proj_num);
+            trk_arr[1]= trk_full[i][0] - eta_ctr;
+            trk_arr[2]= trk_full[i][1] - phi_ctr;
+            trk_arr[3]= trk_full[i][2];
+            trk_arr[4]= 1.;
+
+            cells.emplace_back(trk_arr);
+        }
+    }
+
+    int index = 0;
+    for(auto &in : cells){
+      std::copy(in.begin(), in.end(), inputnn.begin() + index);
+      index+=5;
+      if(index >= static_cast<int>(inputnn.size()-4)) {
+        ATH_MSG_WARNING("Data exceeded tensor size");
+        break;
+      }
+    }
+        
+    //Normalization prior to training
+    NormalizeTensor(inputnn, cells.size() * 5 );
+        
+    float predictedEnergy = exp(runOnnxInference(inputnn)) * 1000.0;//Correct to MeV units 
+    ATH_MSG_DEBUG("NN Predicted energy " << predictedEnergy);      
+    return predictedEnergy;
+
+}
+
+void PFEnergyPredictorTool::NormalizeTensor(std::vector<float> &inputnn, size_t limit) const{
+    size_t i=0;
+    for(i =0;i<limit;i+=5){
+      auto &f = inputnn[i+3];
+      if(f!= 0.0f) f/= 3630.f;
+      auto &e = inputnn[i+0];
+      if(e!= 0.0f){
+       e = std::log(e);
+       e = (e - m_cellE_mean)/m_cellE_std;
+      }
+      auto &eta = inputnn[i+1];
+      if(eta!= 0.0) eta /= 0.7f;
+      auto &phi = inputnn[i+2];
+      if(phi!= 0.0) phi /= m_cellPhi_std;
+    }
+    if(i> inputnn.size()){
+        ATH_MSG_ERROR("Index exceeded tensor MEMORY CORRUPTION");
+    }
+}
+
+
+
+StatusCode PFEnergyPredictorTool::finalize()
+{
+  return StatusCode::SUCCESS;
+}
+

Reconstruction/eflowRec/src/PFSubtractionTool.cxx

@@ -39,6 +39,8 @@ StatusCode PFSubtractionTool::initialize()
   ATH_CHECK(m_theMatchingToolForPull_015.retrieve());
   ATH_CHECK(m_theMatchingToolForPull_02.retrieve());
 
+  if (!m_NNEnergyPredictorTool.empty()) ATH_CHECK(m_NNEnergyPredictorTool.retrieve());
+
   //Set the level of the helpers to the same as the tool here
   m_pfCalc.msg().setLevel(this->msg().level());
   m_pfSubtractionStatusSetter.msg().setLevel(this->msg().level());
@@ -196,7 +198,7 @@ unsigned int PFSubtractionTool::matchAndCreateEflowCaloObj(PFData &data) const{
   //For each eflowCaloObject we calculate the expected energy deposit in the calorimeter and cell ordering for subtraction.  
   for (unsigned int iCalo = nCaloObj; iCalo < data.caloObjects->size(); ++iCalo) {  
     eflowCaloObject* thisEflowCaloObject = data.caloObjects->at(iCalo);
-    thisEflowCaloObject->simulateShower(&integrator, m_binnedParameters.get(), true);        
+    thisEflowCaloObject->simulateShower(&integrator, m_binnedParameters.get(), true, m_useNNEnergy ? &(*m_NNEnergyPredictorTool) : nullptr);        
   }
 
   if (!m_recoverSplitShowers) return nMatches;

Reconstruction/eflowRec/src/components/eflowRec_entries.cxx

@@ -24,6 +24,7 @@
 #include "../PFTrackPreselAlg.h"
 #include "../PFTrackMuonCaloTaggingAlg.h"
 #include "../PFTrackMuonIsoTaggingAlg.h"
+#include "eflowRec/PFEnergyPredictorTool.h"
 
 DECLARE_COMPONENT( PFLeptonSelector )
 DECLARE_COMPONENT( PFClusterSelectorTool )
@@ -47,3 +48,4 @@ DECLARE_COMPONENT( PFTauFlowElementAssoc )
 DECLARE_COMPONENT( PFTrackPreselAlg )
 DECLARE_COMPONENT( PFTrackMuonCaloTaggingAlg )
 DECLARE_COMPONENT( PFTrackMuonIsoTaggingAlg )
+DECLARE_COMPONENT( PFEnergyPredictorTool )
\ No newline at end of file

Reconstruction/eflowRec/src/eflowCaloObject.cxx

@@ -19,6 +19,7 @@ CREATED:  22nd November, 2004
 #include "eflowRec/eflowLayerIntegrator.h"
 #include "eflowRec/eflowEEtaBinnedParameters.h"
 #include "eflowRec/eflowRingSubtractionManager.h"
+#include "eflowRec/PFEnergyPredictorTool.h"
 
 eflowCaloObject::~eflowCaloObject() = default;
 
@@ -60,7 +61,8 @@ double eflowCaloObject::getClusterEnergy() const {
   return clusterEnergy;
 }
 
-void eflowCaloObject::simulateShower(eflowLayerIntegrator *integrator, const eflowEEtaBinnedParameters* binnedParameters, bool useUpdated2015ChargedShowerSubtraction){
+void eflowCaloObject::simulateShower(eflowLayerIntegrator *integrator, const eflowEEtaBinnedParameters* binnedParameters, bool useUpdated2015ChargedShowerSubtraction,
+const PFEnergyPredictorTool* energyP){
 
   for (auto *thisEfRecTrack : m_eflowRecTracks) {
 
@@ -98,7 +100,7 @@ void eflowCaloObject::simulateShower(eflowLayerIntegrator *integrator, const efl
     cellSubtractionManager.getOrdering(binnedParameters, trackE, trackEM1eta, j1st);
 
     /* Set expected energy in the eflowRecTrack object */
-    const double expectedEnergy = cellSubtractionManager.fudgeMean() * thisEfRecTrack->getTrack()->e();
+    const double expectedEnergy = energyP ? energyP->nnEnergyPrediction(thisEfRecTrack) : cellSubtractionManager.fudgeMean() * thisEfRecTrack->getTrack()->e();     
     const double expectedEnergySigma = fabs(cellSubtractionManager.fudgeStdDev() * thisEfRecTrack->getTrack()->e());
 
     const std::vector<eflowTrackClusterLink*>* bestClusters_015 = thisEfRecTrack->getAlternativeClusterMatches("cone_015");

Reconstruction/eflowRec/src/eflowTrackCaloExtensionTool.cxx

@@ -79,6 +79,7 @@ std::unique_ptr<eflowTrackCaloPoints> eflowTrackCaloExtensionTool::execute(const
 
   /*Create a map to index the TrackParameters at calo (owned by the extension) wrt to layers*/
   std::map<eflowCalo::LAYER, const Trk::TrackParameters*> parametersMap;
+  std::map<CaloCell_ID::CaloSample,const Trk::TrackParameters*> tileParametersMap;
 
   /*get the CaloExtension object*/
   const Trk::CaloExtension * extension = nullptr;
@@ -121,14 +122,25 @@ std::unique_ptr<eflowTrackCaloPoints> eflowTrackCaloExtensionTool::execute(const
       } else if (m_trackParametersIdHelper->isEntryToVolume(clParameter.cIdentifier())) {
         parametersMap[getLayer(&clParameter)] = &clParameter;
       }
+
+      CaloCell_ID::CaloSample caloSample = m_trackParametersIdHelper->caloSample(clParameter.cIdentifier());
+
+      if (tileParametersMap[caloSample] == nullptr){
+        tileParametersMap[caloSample] = &clParameter;
+      } else if (m_trackParametersIdHelper->isEntryToVolume(clParameter.cIdentifier())){
+        tileParametersMap[caloSample] = &clParameter;
+      }
+
     }
+
+
     /*
       parametersMap may have several entries for Tile1,2,3.
       The impact is negligible as the eta/phi of these entries are very similar
       https://its.cern.ch/jira/browse/ATLJETMET-242
     */
     
-    return std::make_unique<eflowTrackCaloPoints>(parametersMap);
+    return std::make_unique<eflowTrackCaloPoints>(parametersMap,tileParametersMap);
 
   }
   else{

Reconstruction/eflowRec/src/eflowTrackCaloPoints.cxx

@@ -44,6 +44,24 @@ eflowTrackCaloPoints::eflowTrackCaloPoints(const std::map<eflowCalo::LAYER, cons
 
 }
 
+eflowTrackCaloPoints::eflowTrackCaloPoints(const std::map<eflowCalo::LAYER, const Trk::TrackParameters*> & trackParameters,
+  std::map<CaloCell_ID::CaloSample,const Trk::TrackParameters*>& tileTrackParamaters) :
+  m_isEM1Barrel(trackParameters.begin()->first == eflowCalo::EMB1){
+     /* Fill etaPhiPositions map */
+    std::map<eflowCalo::LAYER, const Trk::TrackParameters*>::const_iterator itPars = trackParameters.begin();
+    std::map<eflowCalo::LAYER, const Trk::TrackParameters*>::const_iterator endPars = trackParameters.end();
+    m_isEM2Barrel = false;
+    for (; itPars != endPars; ++itPars) {
+      setEtaPhi(itPars->first, parToPosition(itPars->second));
+      if (itPars->first == eflowCalo::EMB2) m_isEM2Barrel = true;
+      m_positions[itPars->first] = parToPosition(itPars->second);
+      m_directions[itPars->first] = parToDirection(itPars->second);
+    }
+
+  for (auto firstTileParam : tileTrackParamaters) setEtaPhiTile(firstTileParam.first,parToPosition(firstTileParam.second));
+
+  }
+
 eflowTrackCaloPoints::~eflowTrackCaloPoints() = default;
 
 void eflowTrackCaloPoints::setEtaPhi(eflowCalo::LAYER lay, const Amg::Vector3D& vec) {
@@ -51,6 +69,11 @@ void eflowTrackCaloPoints::setEtaPhi(eflowCalo::LAYER lay, const Amg::Vector3D&
                                                  : m_defaultEtaPhiPosition;
 }
 
+void eflowTrackCaloPoints::setEtaPhiTile(CaloCell_ID::CaloSample lay, const Amg::Vector3D& vec) {
+  m_tileEtaPhiPositions[lay] = (vec != m_nullVector) ? eflowEtaPhiPosition(vec.eta(), vec.phi())
+                                                 : m_defaultEtaPhiPosition;
+}
+
 //only used in eflowTrackExtrapolatorTool
 void eflowTrackCaloPoints::setEtaPhi(eflowCaloENUM layer, double eta, double phi) {
   m_etaPhiPositions[layer] = eflowEtaPhiPosition(eta, phi);
@@ -65,6 +88,12 @@ const eflowEtaPhiPosition& eflowTrackCaloPoints::getEtaPhiPos(eflowCalo::LAYER l
   return (it == m_etaPhiPositions.end()) ? m_defaultEtaPhiPosition : it->second;
 }
 
+const eflowEtaPhiPosition& eflowTrackCaloPoints::getTileEtaPhiPos(CaloCell_ID::CaloSample  layer) const {
+  std::map< CaloCell_ID::CaloSample , eflowEtaPhiPosition>::const_iterator it = m_tileEtaPhiPositions.find(layer);
+  return (it == m_tileEtaPhiPositions.end()) ? m_defaultEtaPhiPosition : it->second;
+}
+
+
 Amg::Vector3D eflowTrackCaloPoints::getPosition(eflowCalo::LAYER layer) {
   std::map<eflowCalo::LAYER, Amg::Vector3D>::const_iterator it = m_positions.find(layer);
   return (it == m_positions.end()) ? m_nullVector : it->second;