diff --git a/Reconstruction/eflowRec/CMakeLists.txt b/Reconstruction/eflowRec/CMakeLists.txt
index 72913b9fc9adb5205e8ca7db2c30b753196a2c2c..b54105740415c8824f8a9557e32da206c1f4c5e2 100644
--- a/Reconstruction/eflowRec/CMakeLists.txt
+++ b/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} )
diff --git a/Reconstruction/eflowRec/eflowRec/PFEnergyPredictorTool.h b/Reconstruction/eflowRec/eflowRec/PFEnergyPredictorTool.h
new file mode 100644
index 0000000000000000000000000000000000000000..4e4466a46e7abba9db7e937bfc0df76654ad7423
--- /dev/null
+++ b/Reconstruction/eflowRec/eflowRec/PFEnergyPredictorTool.h
@@ -0,0 +1,53 @@
+/*
+ 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
+
diff --git a/Reconstruction/eflowRec/eflowRec/PFSubtractionTool.h b/Reconstruction/eflowRec/eflowRec/PFSubtractionTool.h
index fa806a978cdcf657852e2439a600af6d142f91f8..abfdabeebd5bd0c81f1273ba36b3f8ca5b6e4703 100644
--- a/Reconstruction/eflowRec/eflowRec/PFSubtractionTool.h
+++ b/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
diff --git a/Reconstruction/eflowRec/eflowRec/eflowCaloObject.h b/Reconstruction/eflowRec/eflowRec/eflowCaloObject.h
index 17015c3b5f866c0e0a10edecca9ce299e2711ee5..a891a00e97407378452669b0de88b80735662ed8 100644
--- a/Reconstruction/eflowRec/eflowRec/eflowCaloObject.h
+++ b/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:
diff --git a/Reconstruction/eflowRec/eflowRec/eflowTrackCaloPoints.h b/Reconstruction/eflowRec/eflowRec/eflowTrackCaloPoints.h
index 50d7e20bae9785ba87a59f84822162e1536319ff..2d740468d7965b1b9bf8b687398693f602d25d65 100644
--- a/Reconstruction/eflowRec/eflowRec/eflowTrackCaloPoints.h
+++ b/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) {
diff --git a/Reconstruction/eflowRec/python/PFCfg.py b/Reconstruction/eflowRec/python/PFCfg.py
index 9563cefad2b079e6435fc265c85831ac5678c631..56ea2a3cd9ed3c9c6b547e12418151c8425f3f78 100644
--- a/Reconstruction/eflowRec/python/PFCfg.py
+++ b/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):
diff --git a/Reconstruction/eflowRec/python/PFConfigFlags.py b/Reconstruction/eflowRec/python/PFConfigFlags.py
index 16c9e73a2077d78cb3c3894c60ed4d31c1a8493e..fd7b9f0b042a3353a206834891d34bc490b3e8f6 100644
--- a/Reconstruction/eflowRec/python/PFConfigFlags.py
+++ b/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
diff --git a/Reconstruction/eflowRec/src/PFEnergyPredictorTool.cxx b/Reconstruction/eflowRec/src/PFEnergyPredictorTool.cxx
new file mode 100644
index 0000000000000000000000000000000000000000..e4cdc3de2f466e78648696cf7dc5eae00017f278
--- /dev/null
+++ b/Reconstruction/eflowRec/src/PFEnergyPredictorTool.cxx
@@ -0,0 +1,329 @@
+/*
+ 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;
+}
+
diff --git a/Reconstruction/eflowRec/src/PFSubtractionTool.cxx b/Reconstruction/eflowRec/src/PFSubtractionTool.cxx
index 38b8f24b9b0d624cda723f673d10b916bb5bd9e8..c36e6bcd6cb6bf5e1762acc618dc5798c688e032 100644
--- a/Reconstruction/eflowRec/src/PFSubtractionTool.cxx
+++ b/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;
diff --git a/Reconstruction/eflowRec/src/components/eflowRec_entries.cxx b/Reconstruction/eflowRec/src/components/eflowRec_entries.cxx
index 50ae98064ad21b2aec514466d1b83045e3b2e506..5545aa3f2e76627c9ec3cf5943375de80fc95eea 100644
--- a/Reconstruction/eflowRec/src/components/eflowRec_entries.cxx
+++ b/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
diff --git a/Reconstruction/eflowRec/src/eflowCaloObject.cxx b/Reconstruction/eflowRec/src/eflowCaloObject.cxx
index a502e61ff779c0bd5b1fcfee43a99c67c2f42e60..ae30a011da021f54adc8b891abeca4a360f80114 100644
--- a/Reconstruction/eflowRec/src/eflowCaloObject.cxx
+++ b/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");
diff --git a/Reconstruction/eflowRec/src/eflowTrackCaloExtensionTool.cxx b/Reconstruction/eflowRec/src/eflowTrackCaloExtensionTool.cxx
index 8c8349de7a058272b84f0fa4482fa785e44ff381..6d40ff1dd38965e9f25807e8a6fd0d167666cb6a 100644
--- a/Reconstruction/eflowRec/src/eflowTrackCaloExtensionTool.cxx
+++ b/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{
diff --git a/Reconstruction/eflowRec/src/eflowTrackCaloPoints.cxx b/Reconstruction/eflowRec/src/eflowTrackCaloPoints.cxx
index 364a84794408fb4b9fae1594c2d2b6aa360c1666..4e185c2dc914c7523fb3e4b4ed84a96b33c833af 100644
--- a/Reconstruction/eflowRec/src/eflowTrackCaloPoints.cxx
+++ b/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;