MuonSelectionTool.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef MUONSELECTORTOOLS_MUONSELECTIONTOOL_H
#define MUONSELECTORTOOLS_MUONSELECTIONTOOL_H
 
#include "AsgDataHandles/ReadHandleKey.h"
#include "AsgTools/AsgTool.h"
#include "AsgTools/PropertyWrapper.h"
#include "AsgTools/ToolHandle.h"
#include "MuonAnalysisInterfaces/IMuonSelectionTool.h"
#include "PATCore/IAsgSelectionTool.h"
#include "TF1.h"
#include "TFile.h"
#include "TH2D.h"
#include "TMVA/Reader.h"
#include "TSystem.h"  // Replace with PathResolver
#include "xAODEventInfo/EventInfo.h"
#include "AthOnnxInterfaces/IAthInferenceTool.h"
#include "MVAUtils/BDT.h"
#include <string>
#include <vector>
#include <memory>
#include <mutex>
#include <cstdint>
#include <stdexcept>
 
namespace CP {

    class MuonSelectionTool : public virtual IAsgSelectionTool, public virtual IMuonSelectionTool, public asg::AsgTool {
        ASG_TOOL_CLASS2(MuonSelectionTool, IAsgSelectionTool, CP::IMuonSelectionTool)
 
    public:
        MuonSelectionTool(const std::string& name = "MuonSelection");
 
        virtual ~MuonSelectionTool();


        virtual StatusCode initialize() override;



        virtual const asg::AcceptInfo& getAcceptInfo() const override;

        virtual asg::AcceptData accept(const xAOD::IParticle* p) const override;



        virtual asg::AcceptData accept(const xAOD::Muon& mu) const override;

        virtual void setPassesIDCuts(xAOD::Muon&) const override;

        // void setPassesLowPtEfficiencyCuts( xAOD::Muon& mu ) const;

        virtual void setQuality(xAOD::Muon& mu) const override;

        virtual bool passedIDCuts(const xAOD::Muon&) const override;

        virtual bool passedMuonCuts(const xAOD::Muon&) const override;

        virtual bool passedIDCuts(const xAOD::TrackParticle&) const override;

        virtual bool passedHighPtCuts(const xAOD::Muon&) const override;

        virtual bool passedLowPtEfficiencyCuts(const xAOD::Muon&) const override;
        virtual bool passedLowPtEfficiencyCuts(const xAOD::Muon&, xAOD::Muon::Quality thisMu_quality) const override;

        virtual bool passedErrorCutCB(const xAOD::Muon&) const override;

        virtual bool isBadMuon(const xAOD::Muon&) const override;

        virtual xAOD::Muon::Quality getQuality(const xAOD::Muon& mu) const override;

        virtual bool passedCaloTagQuality(const xAOD::Muon& mu) const override;

        virtual bool passedCaloScore(const xAOD::Muon& mu) const override;

        virtual int getResolutionCategory(const xAOD::Muon&) const override;
        
        std::vector<const xAOD::MuonSegment*> getSegmentsSorted(const xAOD::Muon& mu) const;
        
        bool isBIS78(const float eta, const float phi) const;
        bool isBEE(const float eta, const float phi) const;
        bool isBMG(const float eta, const float phi) const;
        
        float qOverPsignificance(const xAOD::Muon& muon) const;
        float rhoPrime(const xAOD::Muon& muon) const;
        
        struct hitSummary {
            uint8_t nprecisionLayers{0};
            uint8_t nprecisionHoleLayers{0};
            uint8_t nGoodPrecLayers{0};
            uint8_t innerSmallHits{0};
            uint8_t innerLargeHits{0};
            uint8_t middleSmallHits{0};
            uint8_t middleLargeHits{0};
            uint8_t outerSmallHits{0};
            uint8_t outerLargeHits{0};
            uint8_t extendedSmallHits{0};
            uint8_t extendedLargeHits{0};
            uint8_t extendedSmallHoles{0};
            uint8_t isSmallGoodSectors{0};
            uint8_t cscUnspoiledEtaHits{0};
            uint8_t etaLayer1STGCHits{0};
            uint8_t etaLayer2STGCHits{0};
            uint8_t MMHits{0};
        };
        
        template <class P, class T, class S> inline void retrieveSummaryValue(const P& muon, T& value, const S type, bool ignoreMissing = false) const {
            if (!muon.summaryValue(value, type) && !ignoreMissing) {
                ATH_MSG_FATAL(__FILE__ << ":" << __LINE__ << " Failed to retrieve summary value " << type);
                throw std::runtime_error("MuonSelectorTool summary retrieval failed");
            }
        }
        
        void fillSummary(const xAOD::Muon& muon, hitSummary& summary) const;

        virtual float getTightNNScore(const xAOD::Muon& muon) const override;
 
    private:
        bool passedLowPtEfficiencyMVACut(const xAOD::Muon&) const;
        bool passedLowPtEfficiencyMVACutRun3(const xAOD::Muon&) const;

        bool passTight(const xAOD::Muon& mu, float rho, float oneOverPSig) const;

        bool passedBMVmimicCut(const xAOD::Muon&) const;

        asg::AcceptInfo m_acceptInfo;
 
        Gaudi::Property<double> m_maxEta{this, "MaxEta", 2.7, "Maximum eta range to select the muons"};
        Gaudi::Property<int>  m_quality{this, "MuQuality", 1,"Quality to select. Values correspond to 0=Tight, 1=Medium, 2=Loose, 3=VeryLoose (only for debug, not supported), 4=HighPt, 5=LowPtEfficiency"};
        Gaudi::Property<bool> m_toroidOff{this, "ToroidOff", false, "Run the tool in Toroid off setup"};
        Gaudi::Property<bool> m_isRun3{this, "IsRun3Geo", false, "Switch to toggle the run 2 & run 3 geometry cuts. The tool will throw an exception if the run number does not match the expectations later"}; 
        Gaudi::Property<bool> m_excludeNSWFromPrecisionLayers{this, "ExcludeNSWFromPrecisionLayers", false, "Cut on the nprecisionlayers ignoring the NSW"};
        Gaudi::Property<bool> m_recalcPrecisionLayerswNSW{this, "RecalcPrecisionLayerswNSW", true, "Recalculate nprecisionlayers with internal criteria for NSW"};
        Gaudi::Property<bool> m_useCaloScore{this, "UseCaloScore", true, "Use CaloScore for calo-tags in the Loose working point. False will revert back to CaloMuonIDTag (not recommended in Rel.22)"};
        Gaudi::Property<int>  m_caloScoreWP{this, "CaloScoreWP", 4,"Currently there are 4 CaloScore WPs for testing; WP4 is the recommended one"};
        Gaudi::Property<bool> m_useLRT{this, "UseLRT", false, "Enable support for LRT muons"};
 
        // Expert development options
        Gaudi::Property<bool> m_turnOffMomCorr{this, "TurnOffMomCorr", false};
        Gaudi::Property<bool> m_disablePtCuts{this, "DisablePtCuts", false};
        Gaudi::Property<bool> m_developMode{this, "ExpertDevelopMode", false};
        Gaudi::Property<bool> m_TrtCutOff{this, "TrtCutOff", true};
        Gaudi::Property<bool> m_SctCutOff{this, "SctCutOff", false};
        Gaudi::Property<bool> m_PixCutOff{this, "PixCutOff", false};
        Gaudi::Property<bool> m_SiHolesCutOff{this, "SiHolesCutOff", false};
        Gaudi::Property<bool> m_useAllAuthors{this, "UseAllAuthors", true};
        Gaudi::Property<bool> m_use2stationMuonsHighPt{this, "Use2stationMuonsHighPt", true, "for users of high-pT working point to choose whether to include 'safe' 2-station muons"};
        Gaudi::Property<bool> m_useMVALowPt{this, "UseMVALowPt", false, "for users of low-pT working point to choose whether to use MVA and whether to include MuTagIMO muons. Not yet supported in release 22, crrently only for development studies."};
        Gaudi::Property<bool> m_useSegmentTaggedLowPt{this, "UseSegmentTaggedLowPt", false, "Use MVA low-pt WP. In development phase"};
        Gaudi::Property<bool> m_geoOnTheFly{this, "AllowSettingGeometryOnTheFly", false,"avoids crash if run2/run3 geo is wrongly set"};
        Gaudi::Property<bool> m_doBadMuonVetoMimic{this, "DoBadMuonVetoMimic", false}; // switch to cut away the tail of very large smearing in MC to mimic the effect of the bad muon veto for 2-station muons in the high-pT selection
        Gaudi::Property<bool> m_allowComm{this, "AcceptCommChain", true, "Accepts muons from the EMEO chain if set to true"};
        Gaudi::Property<bool> m_useBEEBISInHighPtRun3{this, "UseBEEBISInHighPtRun3", true, "Include BEE and BIS7/8 in high-pT WP in Run3"};        
 
        Gaudi::Property<int> m_expertMode_RunNumber{this, "expertMode_RunNumber", 0, "Expert only option. Overwrites RunNumber"};
        Gaudi::Property<unsigned long long> m_expertMode_EvtNumber{this, "expertMode_EvtNumber", 0, "Expert only option. Overwrites EventNumber"};
    
        SG::ReadHandleKey<xAOD::EventInfo> m_eventInfo{this, "EventInfoContName", "EventInfo", "event info key"};
 
        // MVA configs for low-pT working point. Expert use only!
        Gaudi::Property<std::string> m_MVAreaderFile_EVEN_MuidCB{
            this, "MVAreaderFile_EVEN_MuidCB",
            "MuonSelectorTools/190118_PrelimLowPtMVA/LowPtMVA_Weights/BDTG_9JAN2019_MuidCB_EVEN.weights.xml"};
        Gaudi::Property<std::string> m_MVAreaderFile_ODD_MuidCB{
            this, "MVAreaderFile_ODD_MuidCB",
            "MuonSelectorTools/190118_PrelimLowPtMVA/LowPtMVA_Weights/BDTG_9JAN2019_MuidCB_ODD.weights.xml"};
 
        Gaudi::Property<std::string> m_MVAreaderFile_EVEN_MuGirl{
            this, "MVAreaderFile_EVEN_MuGirl",
            "MuonSelectorTools/190118_PrelimLowPtMVA/LowPtMVA_Weights/BDTG_9JAN2019_MuGirl_EVEN.weights.xml"};
        Gaudi::Property<std::string> m_MVAreaderFile_ODD_MuGirl{
            this, "MVAreaderFile_ODD_MuGirl",
            "MuonSelectorTools/190118_PrelimLowPtMVA/LowPtMVA_Weights/BDTG_9JAN2019_MuGirl_ODD.weights.xml"};
 
        Gaudi::Property<std::string> m_MVAreaderFile_MuTagIMO_etaBin1{
            this, "MVAreaderFile_MuTagIMO_etaBin1", "dev/MuonSelectorTools/181121_MuTagIMO_BDT/BDT_NOV2021_MuTagIMO_etaBin1.weights.xml"};
        Gaudi::Property<std::string> m_MVAreaderFile_MuTagIMO_etaBin2{
            this, "MVAreaderFile_MuTagIMO_etaBin2", "dev/MuonSelectorTools/181121_MuTagIMO_BDT/BDT_NOV2021_MuTagIMO_etaBin2.weights.xml"};
        Gaudi::Property<std::string> m_MVAreaderFile_MuTagIMO_etaBin3{
            this, "MVAreaderFile_MuTagIMO_etaBin3", "dev/MuonSelectorTools/181121_MuTagIMO_BDT/BDT_NOV2021_MuTagIMO_etaBin3.weights.xml"};
 
        // subfolder to load from the calibration db
        Gaudi::Property<std::string> m_calibration_version{this, "CalibrationRelease", "PreRec2016_2016-04-13"};
 
        // possible override for the calibration version
        Gaudi::Property<std::string> m_custom_dir{this, "CustomInputFolder", ""};

        Gaudi::Property<std::string> m_BMVcutFile{this, "BMVcutFile",
                                                  "MuonSelectorTools/180620_BMVmimicCutFunctions/BMVmimicCutFunctions.root"};

        StatusCode getHist(TFile* file, const std::string& histName, std::unique_ptr<TH1>& hist) const;
        //
        std::unique_ptr<TH1> m_tightWP_lowPt_rhoCuts;
        std::unique_ptr<TH1> m_tightWP_lowPt_qOverPCuts;
        std::unique_ptr<TH1> m_tightWP_mediumPt_rhoCuts;
        std::unique_ptr<TH1> m_tightWP_highPt_rhoCuts;
        //
        std::unique_ptr<TF1> m_BMVcutFunction_barrel;
        std::unique_ptr<TF1> m_BMVcutFunction_endcap;
 
        // Need run number (or random run number) to apply period-dependent selections.
        // If selection depends only on data taking year, this can be specified by passing
        // argument needOnlyCorrectYear=true, in which case the random run number decoration
        // from the pile-up reweighting tool is not needed.
        unsigned int getRunNumber(bool needOnlyCorrectYear = false) const;
 
        // TMVA readers for low-pT working point
        std::unique_ptr<TMVA::Reader> m_readerE_MUID{nullptr};
        std::unique_ptr<TMVA::Reader> m_readerO_MUID{nullptr};
        std::unique_ptr<TMVA::Reader> m_readerE_MUGIRL{nullptr};
        std::unique_ptr<TMVA::Reader> m_readerO_MUGIRL{nullptr};
 
        std::unique_ptr<TMVA::Reader> m_reader_MUTAGIMO_etaBin1{nullptr};
        std::unique_ptr<TMVA::Reader> m_reader_MUTAGIMO_etaBin2{nullptr};
        std::unique_ptr<TMVA::Reader> m_reader_MUTAGIMO_etaBin3{nullptr};
 
        // MVAUtils for Run-3 low-pT working point
        std::unique_ptr<MVAUtils::BDT> m_MuidCO{nullptr};
        std::vector<double> m_lowPtMuidCO_means;
        std::vector<double> m_lowPtMuidCO_scaler;
        std::unique_ptr<MVAUtils::BDT> m_MuGirl{nullptr};
        std::vector<double> m_lowPtMuGirl_means;
        std::vector<double> m_lowPtMuGirl_scaler;
 
        // variables for the TMVA readers
        mutable std::mutex m_low_pt_mva_mutex;
 
        void IdMsPt(const xAOD::Muon& muon, float& idPt, float& msPt) const;
        
        void checkSanity() const;
 
        bool isRun3(bool forceOnTheFly = false) const {
          if (!m_geoOnTheFly && !forceOnTheFly) return m_isRun3;
          int rn = getRunNumber(true);
          return rn>=399999;
        }
        
        void retrieveParam(const xAOD::Muon& muon, float& value, const xAOD::Muon::ParamDef param) const;
 
        //ONNX
        Gaudi::Property<bool> m_calculateTightNNScore{this, "CalculateTightNNScore", false};
        ToolHandle< AthInfer::IAthInferenceTool >  m_onnxTool{this, "ORTInferenceTool", "AthOnnx::OnnxRuntimeInferenceTool"};
 
    };  // class MuonSelectionTool
 
}  // namespace CP
 
#endif  // CPTOOLTESTS_MUONSELECTIONTOOL_H