diff --git a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CMakeLists.txt b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CMakeLists.txt index 6ed5abefd201a5b0487f457341dafa5ea1717981..474192012ce0870f03af4517af1334e6e9d0e98e 100644 --- a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CMakeLists.txt +++ b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CMakeLists.txt @@ -46,4 +46,4 @@ atlas_add_component( CaloTrkMuIdTools atlas_install_headers( CaloTrkMuIdTools ) atlas_install_joboptions( share/*.py ) atlas_install_runtime( share/CaloMuonLikelihood.PDF.A0.root share/CaloMuonLikelihood.PDF.A1.root share/CaloMuonLikelihood.PDF.A2.root share/CaloMuonLikelihood.PDF.B0.root share/CaloMuonLikelihood.PDF.B1.root share/CaloMuonLikelihood.PDF.B2.root share/CaloMuonLikelihood.PDF.C0.root share/CaloMuonLikelihood.PDF.C1.root share/CaloMuonLikelihood.PDF.C2.root share/CaloTag.CutConfig.root ) - +atlas_install_runtime( share/CaloMuonScoreModels/*.onnx ) diff --git a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CaloTrkMuIdTools/CaloMuonScoreTool.h b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CaloTrkMuIdTools/CaloMuonScoreTool.h new file mode 100644 index 0000000000000000000000000000000000000000..dc8648c33d75a7b95cb25189c4d81a22c482159f --- /dev/null +++ b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/CaloTrkMuIdTools/CaloMuonScoreTool.h @@ -0,0 +1,80 @@ +/* + Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration +*/ + +#ifndef CALOTRKMUIDTOOLS_CALOMUONSCORETOOL_H +#define CALOTRKMUIDTOOLS_CALOMUONSCORETOOL_H + +#include "ICaloTrkMuIdTools/ICaloMuonLikelihoodTool.h" + +#include "AthenaBaseComps/AthAlgTool.h" +#include "GaudiKernel/ToolHandle.h" + +#include "RecoToolInterfaces/IParticleCaloExtensionTool.h" +#include "CaloEvent/CaloClusterContainer.h" + +#include <vector> + + +/** @class CaloMuonScoreTool + + Fetch the calorimeter cells around a track particle and compute the muon score. + + The muon score is computed by doing inference on a 7-colour-channel convolutional neural + network. The inputs to the convolutional neural network are the energy deposits in 30 eta + and 30 phi bins around the track particle. Seven colour channels are considered, + corresponding to the seven calorimeter layers (CaloSamplingIDs) in the low-eta region + (eta < 0.1). + + The convolutional neural network was trained using tensorflow. Inference on this model + is done using ONNX (the tensorflow model having been converted to ONNX format). + + @author ricardo.wolker@cern.ch +*/ +class CaloMuonScoreTool : public AthAlgTool, virtual public ICaloMuonScoreTool { +public: + CaloMuonScoreTool(const std::string& type, const std::string& name, const IInterface* parent); + virtual ~CaloMuonScoreTool()=default; + + virtual StatusCode initialize(); + + // Compute the muon score given a track particle + double getMuonScore(const xAOD::TrackParticle* trk) const; + + // Compute the median of a vector of floats (can be even or odd in length) + float getMedian(std::vector<float> v) const; + + // Get a linearly spaced vector of size `nBins`, ranging from `min` to `max` (both values included) + std::vector<float> getLinearlySpacedBins(float min, float max, int nNins) const; + + // Given a vector of bins, return the index of the matching bin + int getBin(std::vector<float> &bins, float &val) const; + + // Given a calo sampling ID (as integer), return the corresponding "RGB"-like channel ID (0,1,2,3,4,5,6) + int channelForSamplingId(int &samplingId) const; + + +private: + // model name to be loaded + const char MODEL_NAME[] = "./models/model.onnx"; + + // Number of bins in eta + const int ETA_BINS = 30; + + // Number of bins in phi + const int PHI_BINS = 30; + + // window in terms of abs(eta) to consider around the median eta value + const float ETA_CUT = 0.25; + + // window in terms of abs(phi) to consider around the median phi value + const float PHI_CUT = 0.25; + + // Number of colour channels to consider in the convolutional neural network + const int N_CHANNELS = 7; + + + ToolHandle <Trk::IParticleCaloExtensionTool> m_caloExtensionTool{this, "ParticleCaloExtensionTool", ""}; +}; + +#endif diff --git a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloMuonScoreModels/CaloMuonCNN_0.onnx b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloMuonScoreModels/CaloMuonCNN_0.onnx new file mode 100644 index 0000000000000000000000000000000000000000..41c4ea7706a3d03a918527cc119326180876df9f Binary files /dev/null and b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloMuonScoreModels/CaloMuonCNN_0.onnx differ diff --git a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloTrkMuIdTools_jobOptions.py b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloTrkMuIdTools_jobOptions.py index dbd17150fb4cf4887ea50ebf5257a2966672e53c..fddeef89261af8d562e6cf8766da818b79753521 100755 --- a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloTrkMuIdTools_jobOptions.py +++ b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/share/CaloTrkMuIdTools_jobOptions.py @@ -60,3 +60,8 @@ from CaloTrkMuIdTools.CaloTrkMuIdToolsConf import CaloMuonLikelihoodTool as Conf CaloMuonLikelihoodTool = ConfiguredCaloMuonLikelihoodTool(TrackEnergyInCaloTool = TrackEnergyInCaloTool) ToolSvc += CaloMuonLikelihoodTool +### Configure CaloMuonScoreTool ### +from CaloTrkMuIdTools.CaloTrkMuIdToolsConf import CaloMuonScoreTool as ConfiguredCaloMuonScoreTool +CaloMuonScoreTool = ConfiguredCaloMuonScoreTool(TrackEnergyInCaloTool = TrackEnergyInCaloTool) +ToolSvc += CaloMuonScoreTool + diff --git a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/CaloMuonScoreTool.cxx b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/CaloMuonScoreTool.cxx new file mode 100644 index 0000000000000000000000000000000000000000..044c3e0a521badb5d074c12f71fd9da439faee30 --- /dev/null +++ b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/CaloMuonScoreTool.cxx @@ -0,0 +1,80 @@ +/* + Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration +*/ + +#include "CaloTrkMuIdTools/CaloMuonScoreTool.h" + +#include "xAODCaloEvent/CaloCluster.h" +#include "PathResolver/PathResolver.h" +#include "CaloIdentifier/CaloCell_ID.h" +#include "GaudiKernel/SystemOfUnits.h" + +#include "TrkCaloExtension/CaloExtension.h" +#include "TrkCaloExtension/CaloExtensionHelpers.h" +#include "TrkParameters/TrackParameters.h" + +#include "TFile.h" +#include "TH1F.h" + +#include <string> +#include <iostream> +#include <cmath> +#include <map> + +/////////////////////////////////////////////////////////////////////////////// +// CaloMuonScoreTool constructor +/////////////////////////////////////////////////////////////////////////////// +CaloMuonScoreTool::CaloMuonScoreTool(const std::string& type, const std::string& name, const IInterface* parent) : + AthAlgTool(type,name,parent), + m_modelFileNames{"CaloMuonCNN_0.onnx"}; +{ + declareInterface<ICaloMuonScoreTool>(this); + declareProperty("ModelFileNames", m_modelFileNames); +} + +/////////////////////////////////////////////////////////////////////////////// +// CaloMuonScoreTool::initialize +/////////////////////////////////////////////////////////////////////////////// +StatusCode CaloMuonScoreTool::initialize() { + ATH_MSG_INFO("Initializing " << name()); + + ATH_CHECK(m_caloExtensionTool.retrieve()); + + if (m_modelFileNames.size()<1) { + ATH_MSG_FATAL("Number of input ONNX model files should be at least 1"); + return StatusCode::FAILURE; + } + + return StatusCode::SUCCESS; +} + +/////////////////////////////////////////////////////////////////////////////// +// CaloMuonScoreTool::getMuonScore +/////////////////////////////////////////////////////////////////////////////// +double CaloMuonScoreTool::getLHR( const xAOD::TrackParticle* trk, const xAOD::CaloClusterContainer* ClusContainer, const double dR_CUT ) const { + ATH_MSG_DEBUG("in CaloMuonScoreTool::getLHR()"); + + if(trk && ClusContainer){ + double eta_trk = trk->eta(); + double p_trk = 0; + double qOverP = trk->qOverP(); + if (qOverP!=0) p_trk = std::abs(1/qOverP); + + std::unique_ptr<Trk::CaloExtension> caloExt=m_caloExtensionTool->caloExtension(*trk); + + if(!caloExt) return 0; + const Trk::TrackParameters* caloEntryInterSec = caloExt->caloEntryLayerIntersection(); + if(!caloEntryInterSec) { + ATH_MSG_WARNING("getLHR() - caloEntryLayerIntersection is nullptr"); + return 0; + } + double eta_trkAtCalo = caloEntryInterSec->eta(); + double phi_trkAtCalo = caloEntryInterSec->parameters()[Trk::phi]; + + double LR = getLHR( ClusContainer, eta_trk, p_trk, eta_trkAtCalo, phi_trkAtCalo, dR_CUT); + return LR; + } + + return 0; +} + diff --git a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/components/CaloTrkMuIdTools_entries.cxx b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/components/CaloTrkMuIdTools_entries.cxx index 6d46c86e9ddcd8349ea4f45d2b2ef71252f5d79e..d17d5566a3f5dd5495510f0ebfaebc409d973235 100644 --- a/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/components/CaloTrkMuIdTools_entries.cxx +++ b/Reconstruction/MuonIdentification/CaloTrkMuIdTools/src/components/CaloTrkMuIdTools_entries.cxx @@ -2,9 +2,11 @@ #include "CaloTrkMuIdTools/TrackEnergyInCaloTool.h" #include "CaloTrkMuIdTools/TrackDepositInCaloTool.h" #include "CaloTrkMuIdTools/CaloMuonLikelihoodTool.h" +#include "CaloTrkMuIdTools/CaloMuonScoreTool.h" DECLARE_COMPONENT( CaloMuonTag ) DECLARE_COMPONENT( TrackEnergyInCaloTool ) DECLARE_COMPONENT( TrackDepositInCaloTool ) DECLARE_COMPONENT( CaloMuonLikelihoodTool ) +DECLARE_COMPONENT( CaloMuonScoreTool ) diff --git a/Reconstruction/MuonIdentification/ICaloTrkMuIdTools/ICaloTrkMuIdTools/ICaloMuonScoreTool.h b/Reconstruction/MuonIdentification/ICaloTrkMuIdTools/ICaloTrkMuIdTools/ICaloMuonScoreTool.h new file mode 100644 index 0000000000000000000000000000000000000000..38c2c9045e7026db4e7790fee608c290386d3222 --- /dev/null +++ b/Reconstruction/MuonIdentification/ICaloTrkMuIdTools/ICaloTrkMuIdTools/ICaloMuonScoreTool.h @@ -0,0 +1,26 @@ +/* + Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration +*/ + +#ifndef CALOTRKMUIDTOOLS_ICALOMUONSCORETOOL_H +#define CALOTRKMUIDTOOLS_ICALOMUONSCORETOOL_H + +#include "GaudiKernel/IAlgTool.h" +#include "xAODTracking/TrackParticle.h" +#include "xAODCaloEvent/CaloClusterContainer.h" + +static const InterfaceID IID_ICaloMuonScoreTool("ICaloMuonScoreTool",1,0); + +class ICaloMuonScoreTool : virtual public IAlgTool +{ + public: + + virtual ~ICaloMuonScoreTool(){} + + static const InterfaceID& interfaceID(){return IID_ICaloMuonScoreTool;} + + virtual double getMuonScore(const xAOD::TrackParticle* trk) const = 0; + +}; + +#endif diff --git a/Reconstruction/MuonIdentification/MuonCombinedRecExample/python/MuonCaloTagTool.py b/Reconstruction/MuonIdentification/MuonCombinedRecExample/python/MuonCaloTagTool.py index 30842320cc1ed19dbddf27d68d3a0e413f60db3b..4a508067a7e426df99c020a84b4b4f843d7910f7 100644 --- a/Reconstruction/MuonIdentification/MuonCombinedRecExample/python/MuonCaloTagTool.py +++ b/Reconstruction/MuonIdentification/MuonCombinedRecExample/python/MuonCaloTagTool.py @@ -47,6 +47,10 @@ def CaloMuonLikelihoodTool(name='CaloMuonLikelihoodTool', **kwargs ): kwargs.setdefault("ParticleCaloExtensionTool", getPublicTool("MuonParticleCaloExtensionTool") ) return CfgMgr.CaloMuonLikelihoodTool(name,**kwargs) +def CaloMuonScoreTool(name='CaloMuonScoreTool', **kwargs ): + kwargs.setdefault("ParticleCaloExtensionTool", getPublicTool("MuonParticleCaloExtensionTool") ) + return CfgMgr.CaloMuonScoreTool(name,**kwargs) + def MuonCaloTagTool( name='MuonCaloTagTool', **kwargs ): from CaloTrkMuIdTools.CaloTrkMuIdToolsConf import CaloMuonTag as ConfiguredCaloMuonTag CaloMuonTagLoose = ConfiguredCaloMuonTag(name = "CaloMuonTagLoose")