MuonAnalysisConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
 
# AnaAlgorithm import(s):
from AnalysisAlgorithmsConfig.ConfigBlock import ConfigBlock
from AnalysisAlgorithmsConfig.ConfigSequence import groupBlocks
from AthenaCommon.SystemOfUnits import GeV
from AnalysisAlgorithmsConfig.ConfigAccumulator import DataType
from TrackingAnalysisAlgorithms.TrackingAnalysisConfig import InDetTrackCalibrationConfig
from AthenaConfiguration.Enums import LHCPeriod
from TrigGlobalEfficiencyCorrection.TriggerLeg_DictHelpers import TriggerDict
from Campaigns.Utils import Campaign
from AthenaCommon.Logging import logging
 
 
class MuonCalibrationConfig (ConfigBlock):
    """the ConfigBlock for the muon four-momentum correction"""
 
    def __init__ (self) :
        super (MuonCalibrationConfig, self).__init__ ()
        self.setBlockName('Muons')
        self.addOption ('inputContainer', '', type=str,
            info="the name of the input muon container. If left empty, automatically defaults "
            "to `AnalysisMuons` for PHYSLITE or `Muons` otherwise.")
        self.addOption ('containerName', '', type=str,
            noneAction='error',
            info="the name of the output container after calibration.")
        self.addOption ('postfix', "", type=str,
            info="a postfix to apply to decorations and algorithm names. "
            "Typically not needed here since the calibration is common to "
            "all muons.")
        self.addOption ('minPt', 3.0*GeV, type=float,
            info=r"$p_\mathrm{T}$ cut (in MeV) to apply to calibrated muons.")
        self.addOption ('recalibratePhyslite', True, type=bool,
            info="whether to run the `CP::MuonCalibrationAndSmearingAlg` on "
            "PHYSLITE derivations.")
        self.addOption ('maxEta', 2.7, type=float,
            info=r"maximum muon $\vert\eta\vert$.")
        self.addOption ('excludeNSWFromPrecisionLayers', False, type=bool,
            info="only for testing purposes, turn on to ignore NSW hits and "
            "fix a crash with older derivations (p-tag <p5834).")
        self.addOption ('calibMode', 'correctData_CB', type=str, info='calibration mode of the `MuonCalibTool` needed to turn on the sagitta bias corrections and to select the muon track calibration type (CB or ID+MS), see https://atlas-mcp.docs.cern.ch/guidelines/muonmomentumcorrections/index.html#cpmuoncalibtool-tool.')
        self.addOption ('decorateTruth', False, type=bool,
            info="decorate truth particle information on the reconstructed one.")
        self.addOption ('writeTrackD0Z0', False, type = bool,
            info=r"save the $d_0$ significance and $z_0\sin\theta$ variables.")
        self.addOption ('writeColumnarToolVariables', False, type=bool,
            info="whether to add variables needed for running the columnar muon tool(s) on the output n-tuple (EXPERIMENTAL).",
            expertMode=True)
        self.addOption ('runTrackBiasing', False, type=bool,
            info="EXPERIMENTAL: This enables the `InDetTrackBiasingTool`, for tracks "
            "associated to muons. The tool does not have Run 3 recommendations yet.",
            expertMode=True)
        self.addOption ('addGlobalFELinksDep', False, type=bool,
            info="whether to add dependencies for the global FE links (needed for PHYSLITE production)",
            expertMode=True)
 
    def instanceName (self) :
        if self.postfix != "":
            return self.postfix
        else :
            return self.containerName
 
    def makeAlgs (self, config) :
 
        log = logging.getLogger('MuonCalibrationConfig')
 
        #make sure that this is sync with
        #PhysicsAnalysis/MuonID/MuonIDAnalysis/MuonMomentumCorrections/MuonMomentumCorrections/MuonCalibTool.h#L31-37
        if self.calibMode == 'correctData_CB':
            calibMode = 0
        elif self.calibMode == 'correctData_IDMS':
            calibMode = 1
        elif self.calibMode == 'notCorrectData_IDMS':
            calibMode = 2
        elif self.calibMode == 'notCorrectData_CB':
            calibMode = 3
        else :
            raise ValueError ("invalid calibMode: \"" + self.calibMode + "\". Allowed values are correctData_CB, correctData_IDMS, notCorrectData_IDMS, notCorrectData_CB")
 
        inputContainer = "AnalysisMuons" if config.isPhyslite() else "Muons"
        if self.inputContainer:
            inputContainer = self.inputContainer
        config.setSourceName (self.containerName, inputContainer)
        config.setContainerMeta (self.containerName, 'calibMode', calibMode)
 
        # Set up a shallow copy to decorate
        if config.wantCopy (self.containerName) :
            alg = config.createAlgorithm( 'CP::AsgShallowCopyAlg', 'MuonShallowCopyAlg' )
            alg.input = config.readName (self.containerName)
            alg.output = config.copyName (self.containerName)
            alg.outputType = 'xAOD::MuonContainer'
            decorationList = ['DFCommonJetDr','neflowisol20_CloseByCorr',
                              'ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt1000_CloseByCorr',
                              'ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt500_CloseByCorr',
                              'topoetcone20_CloseByCorr']
            if self.addGlobalFELinksDep:
                decorationList += ['neutralGlobalFELinks', 'chargedGlobalFELinks']
            if config.dataType() is not DataType.Data:
                decorationList += ['TruthLink']
            if self.addGlobalFELinksDep or config.dataType() is not DataType.Data:
                alg.declareDecorations = decorationList
 
        # Set up the eta-cut on all muons prior to everything else
        alg = config.createAlgorithm( 'CP::AsgSelectionAlg',
                               'MuonEtaCutAlg' )
        config.addPrivateTool( 'selectionTool', 'CP::AsgPtEtaSelectionTool' )
        alg.selectionTool.maxEta = self.maxEta
        alg.selectionDecoration = 'selectEta' + self.postfix + ',as_bits'
        alg.particles = config.readName (self.containerName)
        alg.preselection = config.getPreselection (self.containerName, '')
        config.addSelection (self.containerName, '', alg.selectionDecoration)
 
        # Set up the muon calibration and smearing algorithm:
        alg = config.createAlgorithm( 'CP::MuonCalibrationAndSmearingAlg',
                               'MuonCalibrationAndSmearingAlg' )
        config.addPrivateTool( 'calibrationAndSmearingTool',
                        'CP::MuonCalibTool' )
 
        alg.calibrationAndSmearingTool.IsRun3Geo = config.geometry() >= LHCPeriod.Run3
        alg.calibrationAndSmearingTool.calibMode = calibMode
        if config.geometry() is LHCPeriod.Run4:
            log.warning("Disabling NSW hits for Run4 geometry")
            alg.calibrationAndSmearingTool.ExcludeNSWFromPrecisionLayers = True
        else:
            alg.calibrationAndSmearingTool.ExcludeNSWFromPrecisionLayers = self.excludeNSWFromPrecisionLayers and (config.geometry() >= LHCPeriod.Run3)
        alg.muons = config.readName (self.containerName)
        alg.muonsOut = config.copyName (self.containerName)
        alg.preselection = config.getPreselection (self.containerName, '')
        if config.isPhyslite() and not self.recalibratePhyslite :
            alg.skipNominal = True
 
        # Set up the the pt selection
        if self.minPt > 0:
            alg = config.createAlgorithm( 'CP::AsgSelectionAlg', 'MuonPtCutAlg' )
            alg.selectionDecoration = 'selectPt' + self.postfix + ',as_bits'
            config.addPrivateTool( 'selectionTool', 'CP::AsgPtEtaSelectionTool' )
            alg.particles = config.readName (self.containerName)
            alg.selectionTool.minPt = self.minPt
            alg.preselection = config.getPreselection (self.containerName, '')
            config.addSelection (self.containerName, '', alg.selectionDecoration,
                                preselection = True)
 
        # Additional decorations
        if self.writeTrackD0Z0:
            alg = config.createAlgorithm( 'CP::AsgLeptonTrackDecorationAlg',
                                          'LeptonTrackDecorator' )
            if config.dataType() is not DataType.Data:
                if self.runTrackBiasing:
                    InDetTrackCalibrationConfig.makeTrackBiasingTool(config, alg)
                InDetTrackCalibrationConfig.makeTrackSmearingTool(config, alg)
            alg.particles = config.readName (self.containerName)
 
        alg = config.createAlgorithm( 'CP::AsgEnergyDecoratorAlg', 'EnergyDecorator' )
        alg.particles = config.readName (self.containerName)
 
        config.addOutputVar (self.containerName, 'pt', 'pt')
        config.addOutputVar (self.containerName, 'eta', 'eta', noSys=True)
        config.addOutputVar (self.containerName, 'phi', 'phi', noSys=True)
        config.addOutputVar (self.containerName, 'e_%SYS%', 'e')
        config.addOutputVar (self.containerName, 'charge', 'charge', noSys=True)
 
        if self.writeTrackD0Z0:
            config.addOutputVar (self.containerName, 'd0_%SYS%', 'd0')
            config.addOutputVar (self.containerName, 'd0sig_%SYS%', 'd0sig')
            config.addOutputVar (self.containerName, 'z0_%SYS%', 'z0')
            config.addOutputVar (self.containerName, 'z0sintheta_%SYS%', 'z0sintheta')
            config.addOutputVar (self.containerName, 'z0sinthetasig_%SYS%', 'z0sinthetasig')
 
        # decorate truth information on the reconstructed object:
        if self.decorateTruth and config.dataType() is not DataType.Data:
            config.addOutputVar (self.containerName, "truthType", "truth_type", noSys=True)
            config.addOutputVar (self.containerName, "truthOrigin", "truth_origin", noSys=True)
        
        config.addOutputVar (self.containerName, 'muonType', 'muonType', noSys=True, enabled=self.writeColumnarToolVariables)
 
class MuonWorkingPointSelectionConfig (ConfigBlock) :
    """the ConfigBlock for the muon working point selection"""
 
    def __init__ (self) :
        super (MuonWorkingPointSelectionConfig, self).__init__ ()
        self.setBlockName('MuonWorkingPointSelection')
        self.addOption ('containerName', '', type=str,
            noneAction='error',
            info="the name of the input container.")
        self.addOption ('selectionName', '', type=str,
            noneAction='error',
            info="the name of the muon selection to define (e.g. `tight` or `loose`).")
        self.addOption ('postfix', None, type=str,
            info="a postfix to apply to decorations and algorithm names. "
            "Typically not needed here as `selectionName` is used internally.")
        self.addOption ('trackSelection', True, type=bool,
            info="whether or not to set up an instance of "
            "`CP::AsgLeptonTrackSelectionAlg`, with the recommended $d_0$ and "
            r"$z_0\sin\theta$ cuts.")
        self.addOption ('maxD0Significance', 3, type=float,
            info="maximum $d_0$ significance used for the track selection.")
        self.addOption ('maxDeltaZ0SinTheta', 0.5, type=float,
            info=r"maximum $\Delta z_0\sin\theta$ (in mm) used for the track selection.")
        self.addOption ('quality', None, type=str,
            info="the ID WP to use. Supported ID WPs: `Tight`, `Medium`, "
            "`Loose`, `LowPt`, `HighPt`.")
        self.addOption ('isolation', None, type=str,
            info="the isolation WP to use. Supported isolation WPs: "
            "`PflowLoose_VarRad`, `PflowTight_VarRad`, `Loose_VarRad`, "
            "`Tight_VarRad`, `NonIso`.")
        self.addOption ('addSelectionToPreselection', True, type=bool,
            info="whether to retain only muons satisfying the working point "
            "requirements.")
        self.addOption ('isoDecSuffix', '', type=str,
            info="the `isoDecSuffix` name if using close-by-corrected isolation working points.")
        self.addOption ('excludeNSWFromPrecisionLayers', False, type=bool,
            info="only for testing purposes, turn on to ignore NSW hits and "
            "fix a crash with older derivations (p-tag <p5834).")
    
    def instanceName (self) :
        if self.postfix is not None:
            return self.containerName + '_' + self.postfix
        else:
            return self.containerName + '_' + self.selectionName
 
    def makeAlgs (self, config) :
        log = logging.getLogger('MuonWorkingPointSelectionConfig')
 
        from xAODMuon.xAODMuonEnums import xAODMuonEnums
        if self.quality == 'Tight' :
            quality = xAODMuonEnums.Quality.Tight
        elif self.quality == 'Medium' :
            quality = xAODMuonEnums.Quality.Medium
        elif self.quality == 'Loose' :
            quality = xAODMuonEnums.Quality.Loose
        elif self.quality == 'VeryLoose' :
            quality = xAODMuonEnums.Quality.VeryLoose
        elif self.quality == 'HighPt' :
            quality = 4
        elif self.quality == 'LowPt' :
            quality = 5
        else :
            raise ValueError ("invalid muon quality: \"" + self.quality +
                              "\", allowed values are Tight, Medium, Loose, " +
                              "VeryLoose, HighPt, LowPt")
 
        # The setup below is inappropriate for Run 1
        if config.geometry() is LHCPeriod.Run1:
            raise ValueError ("Can't set up the MuonWorkingPointSelectionConfig with %s, there must be something wrong!" % config.geometry().value)
 
        postfix = self.postfix
        if postfix is None :
            postfix = self.selectionName
        if postfix != '' and postfix[0] != '_' :
            postfix = '_' + postfix
 
        # Set up the track selection algorithm:
        if self.trackSelection:
            alg = config.createAlgorithm( 'CP::AsgLeptonTrackSelectionAlg',
                                          'MuonTrackSelectionAlg',
                                          reentrant=True )
            alg.selectionDecoration = 'trackSelection' + postfix + ',as_bits'
            alg.maxD0Significance = self.maxD0Significance
            alg.maxDeltaZ0SinTheta = self.maxDeltaZ0SinTheta
            alg.particles = config.readName (self.containerName)
            alg.preselection = config.getPreselection (self.containerName, '')
            config.addSelection (self.containerName, self.selectionName, alg.selectionDecoration, preselection=self.addSelectionToPreselection)
 
        # Setup the muon quality selection
        alg = config.createAlgorithm( 'CP::MuonSelectionAlgV2',
                               'MuonSelectionAlg' )
        config.addPrivateTool( 'selectionTool', 'CP::MuonSelectionTool' )
        alg.selectionTool.MuQuality = quality
        alg.selectionTool.IsRun3Geo = config.geometry() >= LHCPeriod.Run3
        if config.geometry() is LHCPeriod.Run4:
            log.warning("Disabling NSW hits for Run4 geometry")
            alg.selectionTool.ExcludeNSWFromPrecisionLayers = True
        else:
            alg.selectionTool.ExcludeNSWFromPrecisionLayers = self.excludeNSWFromPrecisionLayers and (config.geometry() >= LHCPeriod.Run3)
        alg.selectionDecoration = 'good_muon' + postfix + ',as_char'
        alg.badMuonVetoDecoration = 'is_bad' + postfix + ',as_char'
        alg.muons = config.readName (self.containerName)
        alg.preselection = config.getPreselection (self.containerName, self.selectionName)
        config.addSelection (self.containerName, self.selectionName,
                             alg.selectionDecoration,
                             preselection=self.addSelectionToPreselection)
        if self.quality == 'HighPt':
            config.addOutputVar (self.containerName, 'is_bad' + postfix, 'is_bad' + postfix)
 
        # Set up the isolation calculation algorithm:
        if self.isolation != 'NonIso' :
            alg = config.createAlgorithm( 'CP::MuonIsolationAlg',
                                   'MuonIsolationAlg' )
            config.addPrivateTool( 'isolationTool', 'CP::IsolationSelectionTool' )
            alg.isolationTool.MuonWP = self.isolation
            alg.isolationTool.IsoDecSuffix = self.isoDecSuffix
            alg.isolationDecoration = 'isolated_muon' + postfix + ',as_char'
            alg.muons = config.readName (self.containerName)
            alg.preselection = config.getPreselection (self.containerName, self.selectionName)
            config.addSelection (self.containerName, self.selectionName,
                                 alg.isolationDecoration,
                                 preselection=self.addSelectionToPreselection)
 
 
class MuonWorkingPointEfficiencyConfig (ConfigBlock) :
    """the ConfigBlock for the muon working point efficiency computation"""
 
    def __init__ (self) :
        super (MuonWorkingPointEfficiencyConfig, self).__init__ ()
        self.setBlockName('MuonWorkingPointEfficiency')
        self.addDependency('MuonWorkingPointSelection', required=True)
        self.addDependency('EventSelection', required=False)
        self.addDependency('EventSelectionMerger', required=False)
        self.addOption ('containerName', '', type=str,
            noneAction='error',
            info="the name of the input container.")
        self.addOption ('selectionName', '', type=str,
            noneAction='error',
            info="the name of the muon selection to define (e.g. `tight` or `loose`).")
        self.addOption ('postfix', None, type=str,
            info="a postfix to apply to decorations and algorithm names. "
            "Typically not needed here as `selectionName` is used internally.")
        self.addOption ('trackSelection', True, type=bool,
            info="whether or not to set up an instance of "
            "`CP::AsgLeptonTrackSelectionAlg`, with the recommended $d_0$ and "
            r"$z_0\sin\theta$ cuts.")
        self.addOption ('quality', None, type=str,
            info="the ID WP to use. Supported ID WPs: `Tight`, `Medium`, "
            "`Loose`, `LowPt`, `HighPt`.")
        self.addOption ('isolation', None, type=str,
            info="the isolation WP to use. Supported isolation WPs: "
            "`PflowLoose_VarRad`, `PflowTight_VarRad`, `Loose_VarRad`, "
            "`Tight_VarRad`, `NonIso`.")
        self.addOption ('systematicBreakdown', False, type=bool,
            info="enables the full breakdown of efficiency SF systematics "
            "(1 NP per uncertainty source, instead of 1 NP in total).")
        self.addOption ('noEffSF', False, type=bool,
            info="disables the calculation of efficiencies and scale factors. "
            "Experimental! Only useful to test a new WP for which scale "
            "factors are not available.",
            expertMode=True)
        self.addOption ('saveDetailedSF', True, type=bool,
            info="save all the independent detailed object scale factors.")
        self.addOption ('saveCombinedSF', False, type=bool,
            info="save the combined object scale factor.")
    
    def instanceName (self) :
        if self.postfix is not None:
            return self.containerName + '_' + self.postfix
        else:
            return self.containerName + '_' + self.selectionName
 
    def makeAlgs (self, config) :
 
        # The setup below is inappropriate for Run 1
        if config.geometry() is LHCPeriod.Run1:
            raise ValueError ("Can't set up the MuonWorkingPointEfficiencyConfig with %s, there must be something wrong!" % config.geometry().value)
 
        postfix = self.postfix
        if postfix is None :
            postfix = self.selectionName
        if postfix != '' and postfix[0] != '_' :
            postfix = '_' + postfix
 
        sfList = []
        # Set up the reco/ID efficiency scale factor calculation algorithm:
        if config.dataType() is not DataType.Data and not self.noEffSF:
            alg = config.createAlgorithm( 'CP::MuonEfficiencyScaleFactorAlg',
                                   'MuonEfficiencyScaleFactorAlgReco' )
            config.addPrivateTool( 'efficiencyScaleFactorTool',
                            'CP::MuonEfficiencyScaleFactors' )
            config.setExtraInputs ({('xAOD::EventInfo', 'EventInfo.RandomRunNumber')})
            alg.scaleFactorDecoration = 'muon_reco_effSF' + postfix + "_%SYS%"
            alg.outOfValidity = 2 #silent
            alg.outOfValidityDeco = 'muon_reco_bad_eff' + postfix
            alg.efficiencyScaleFactorTool.WorkingPoint = self.quality
            if config.geometry() >= LHCPeriod.Run3:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '251211_Preliminary_r24run3'
            else:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '230213_Preliminary_r22run2'
            alg.efficiencyScaleFactorTool.BreakDownSystematics = self.systematicBreakdown
            alg.muons = config.readName (self.containerName)
            alg.preselection = config.getPreselection (self.containerName, self.selectionName)
            if self.saveDetailedSF:
                config.addOutputVar (self.containerName, alg.scaleFactorDecoration,
                                     'reco_effSF' + postfix)
            sfList += [alg.scaleFactorDecoration]
 
        # Set up the HighPt-specific BadMuonVeto efficiency scale factor calculation algorithm:
        if config.dataType() is not DataType.Data and self.quality == 'HighPt' and not self.noEffSF:
            alg = config.createAlgorithm( 'CP::MuonEfficiencyScaleFactorAlg',
                                   'MuonEfficiencyScaleFactorAlgBMVHighPt' )
            config.addPrivateTool( 'efficiencyScaleFactorTool',
                            'CP::MuonEfficiencyScaleFactors' )
            alg.scaleFactorDecoration = 'muon_BadMuonVeto_effSF' + postfix + "_%SYS%"
            alg.outOfValidity = 2 #silent
            alg.outOfValidityDeco = 'muon_BadMuonVeto_bad_eff' + postfix
            alg.efficiencyScaleFactorTool.WorkingPoint = 'BadMuonVeto_HighPt'
            if config.geometry() >= LHCPeriod.Run3:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '220817_Preliminary_r22run3' # not available as part of '230123_Preliminary_r22run3'!
            else:
                 alg.efficiencyScaleFactorTool.CalibrationRelease = '230213_Preliminary_r22run2'
            alg.efficiencyScaleFactorTool.BreakDownSystematics = self.systematicBreakdown
            alg.muons = config.readName (self.containerName)
            alg.preselection = config.getPreselection (self.containerName, self.selectionName)
            if self.saveDetailedSF:
                config.addOutputVar (self.containerName, alg.scaleFactorDecoration,
                                     'BadMuonVeto_effSF' + postfix)
            sfList += [alg.scaleFactorDecoration]
 
        # Set up the isolation efficiency scale factor calculation algorithm:
        if config.dataType() is not DataType.Data and self.isolation != 'NonIso' and not self.noEffSF:
            alg = config.createAlgorithm( 'CP::MuonEfficiencyScaleFactorAlg',
                                   'MuonEfficiencyScaleFactorAlgIsol' )
            config.addPrivateTool( 'efficiencyScaleFactorTool',
                            'CP::MuonEfficiencyScaleFactors' )
            alg.scaleFactorDecoration = 'muon_isol_effSF' + postfix + "_%SYS%"
            alg.outOfValidity = 2 #silent
            alg.outOfValidityDeco = 'muon_isol_bad_eff' + postfix
            alg.efficiencyScaleFactorTool.WorkingPoint = self.isolation + 'Iso'
            if config.geometry() >= LHCPeriod.Run3:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '251211_Preliminary_r24run3'
            else:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '230213_Preliminary_r22run2'
            alg.efficiencyScaleFactorTool.BreakDownSystematics = self.systematicBreakdown
            alg.muons = config.readName (self.containerName)
            alg.preselection = config.getPreselection (self.containerName, self.selectionName)
            if self.saveDetailedSF:
                config.addOutputVar (self.containerName, alg.scaleFactorDecoration,
                                     'isol_effSF' + postfix)
            sfList += [alg.scaleFactorDecoration]
 
        # Set up the TTVA scale factor calculation algorithm:
        if config.dataType() is not DataType.Data and self.trackSelection and not self.noEffSF:
            alg = config.createAlgorithm( 'CP::MuonEfficiencyScaleFactorAlg',
                                   'MuonEfficiencyScaleFactorAlgTTVA' )
            config.addPrivateTool( 'efficiencyScaleFactorTool',
                            'CP::MuonEfficiencyScaleFactors' )
            alg.scaleFactorDecoration = 'muon_TTVA_effSF' + postfix + "_%SYS%"
            alg.outOfValidity = 2 #silent
            alg.outOfValidityDeco = 'muon_TTVA_bad_eff' + postfix
            alg.efficiencyScaleFactorTool.WorkingPoint = 'TTVA'
            if config.geometry() >= LHCPeriod.Run3:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '251211_Preliminary_r24run3'
            else:
                alg.efficiencyScaleFactorTool.CalibrationRelease = '230213_Preliminary_r22run2'
            alg.efficiencyScaleFactorTool.BreakDownSystematics = self.systematicBreakdown
            alg.muons = config.readName (self.containerName)
            alg.preselection = config.getPreselection (self.containerName, self.selectionName)
            if self.saveDetailedSF:
                config.addOutputVar (self.containerName, alg.scaleFactorDecoration,
                                     'TTVA_effSF' + postfix)
            sfList += [alg.scaleFactorDecoration]
 
        if config.dataType() is not DataType.Data and not self.noEffSF and self.saveCombinedSF:
            alg = config.createAlgorithm( 'CP::AsgObjectScaleFactorAlg',
                                          'MuonCombinedEfficiencyScaleFactorAlg' )
            alg.particles = config.readName (self.containerName)
            alg.inScaleFactors = sfList
            alg.outScaleFactor = 'effSF' + postfix + '_%SYS%'
            config.addOutputVar (self.containerName, alg.outScaleFactor, 'effSF' + postfix)
 
class MuonTriggerAnalysisSFBlock (ConfigBlock):
 
    def __init__ (self) :
        super (MuonTriggerAnalysisSFBlock, self).__init__ ()
        self.addDependency('EventSelection', required=False)
        self.addDependency('EventSelectionMerger', required=False)
        self.addOption ('triggerChainsPerYear', {}, type=dict,
                        info="a dictionary with key (string) the year and value (list of "
                        "strings) the trigger chains.")
        self.addOption ('muonID', '', type=str,
                        info="the muon quality WP to use.")
        self.addOption ('saveSF', True, type=bool,
                        info="whether to decorate the trigger scale factor.")
        self.addOption ('saveEff', False, type=bool,
                        info="whether to decorate the trigger MC efficiencies.")
        self.addOption ('saveEffData', False, type=bool,
                        info="whether to decorate the trigger data efficiencies.")
        self.addOption ('prefixSF', 'trigEffSF', type=str,
                        info="the decoration prefix for trigger scale factors.")
        self.addOption ('prefixEff', 'trigEff', type=str,
                        info="the decoration prefix for MC trigger efficiencies.")
        self.addOption ('prefixEffData', 'trigEffData', type=str,
                        info="the decoration prefix for data trigger efficiencies.")
        self.addOption ('includeAllYearsPerRun', False, type=bool,
                        info="all configured years in the LHC run will "
                        "be included in all jobs.")
        self.addOption ('removeHLTPrefix', True, type=bool,
                        info="remove the HLT prefix from trigger chain names.")
        self.addOption ('containerName', '', type=str,
                        info="the input muon container, with a possible selection, in "
                        "the format `container` or `container.selection`.")
        self.addOption ('customToolSuffix', '', type=str,
                        expertMode=True, info="EXPERIMENTAL: specify custom suffix for the public tool name")
        self.addOption ('customInputFolder', '', type=str,
                        expertMode=True, info="EXPERIMENTAL: specify custom input folder")
        self.addOption ('customInputFilePerYear', {}, type=dict,
                        expertMode=True, info="EXPERIMENTAL: specify custom input file per year")
 
    def instanceName (self) :
        return self.containerName + '_' + self.muonID
 
    def makeAlgs (self, config) :
 
        if config.dataType() is not DataType.Data:
 
            # Dictionary from TrigGlobalEfficiencyCorrection/Triggers.cfg
            # Key is trigger chain (w/o HLT prefix)
            # Value is empty for single leg trigger or list of legs
            triggerDict = TriggerDict()
 
            if self.includeAllYearsPerRun:
                years = [int(year) for year in self.triggerChainsPerYear.keys()]
            else:
                from TriggerAnalysisAlgorithms.TriggerAnalysisSFConfig import (
                    get_input_years)
                years = get_input_years(config)
 
            triggerYearStartBoundaries = {
                2015: 260000,
                2016: 290000,
                2017: 324000,
                2018: 348000,
                2022: 410000,
                2023: 450000,
                2024: 470000,
            }
 
            triggerConfigs = {}
            triggerConfigYears = {}
            from TriggerAnalysisAlgorithms.TriggerAnalysisConfig import is_year_in_current_period
            for year in years:
                if not is_year_in_current_period(config, year):
                    continue
 
                triggerChains = self.triggerChainsPerYear.get(int(year), self.triggerChainsPerYear.get(str(year), []))
                for chain in triggerChains:
                    chain = chain.replace(" || ", "_OR_")
                    chain_noHLT = chain.replace("HLT_", "")
                    chain_out = chain_noHLT if self.removeHLTPrefix else chain
                    legs = triggerDict[chain_noHLT]
                    if not legs:
                        if chain_noHLT.startswith('mu') and chain_noHLT[2].isdigit:
                            # Need to support HLT_mu26_ivarmedium_OR_HLT_mu50
                            triggerConfigs[chain_out] = chain
                            if chain_out in triggerConfigYears.keys():
                                triggerConfigYears[chain_out].append(year)
                            else:
                                triggerConfigYears[chain_out] = [year]
                    else:
                        for leg in legs:
                            if leg.startswith('mu') and leg[2].isdigit:
                                # Need to support HLT_mu14_ivarloose
                                leg_out = leg if self.removeHLTPrefix else f"HLT_{leg}"
                                triggerConfigs[leg_out] = f"HLT_{leg}"
                                if leg_out in triggerConfigYears.keys():
                                    triggerConfigYears[leg_out].append(year)
                                else:
                                    triggerConfigYears[leg_out] = [year]
 
            if not triggerConfigs:
                return
 
            # Make the public tool for this configuration
            sfTool = config.createPublicTool("CP::MuonTriggerScaleFactors", f"{self.instanceName()}_SFTool{self.customToolSuffix}")
            # Reproduce config from TrigGlobalEfficiencyAlg
            sfTool.MuonQuality = self.muonID
            sfTool.AllowZeroSF = True
            sfTool.CustomInputFolder = self.customInputFolder
            sfTool.CustomInputFilePerYear = self.customInputFilePerYear
            sfTool.Campaign = config.campaign().value
 
            for trig_short, trig in triggerConfigs.items():
                alg = config.createAlgorithm('CP::MuonTriggerEfficiencyScaleFactorAlg',
                                             'MuonTrigEfficiencyCorrectionsAlg_' + trig_short)
                alg.efficiencyScaleFactorTool = f"{sfTool.getType()}/{sfTool.getName()}"
 
                # Avoid warnings for missing triggers
                if self.includeAllYearsPerRun:
                    alg.minRunNumber = 0
                    alg.maxRunNumber = 999999
 
                    if triggerConfigYears[trig_short][0] != years[0]:
                        alg.minRunNumber = triggerYearStartBoundaries.get(triggerConfigYears[trig_short][0], 999999)
                    if triggerConfigYears[trig_short][-1] != years[-1]:
                        alg.maxRunNumber = triggerYearStartBoundaries.get(triggerConfigYears[trig_short][-1] + 1, 999999)
                elif config.campaign() is Campaign.MC20a:  # to avoid potential corner-cases keep the default config unchanged
                    if triggerConfigYears[trig_short] == [2015]:
                        alg.maxRunNumber = 290000
                    elif triggerConfigYears[trig_short] == [2016]:
                        alg.minRunNumber = 290000
 
                alg.trigger = trig
 
                # Some triggers in `250731_SummerUpdate` recommendations are not supported in 2022 period F
                if config.campaign() is Campaign.MC23a and (trig_short == "HLT_mu8noL1_FSNOSEED" or trig_short == "HLT_mu22_L1MU14FCH"):
                    alg.minRunNumber = 435816  # Start of 2022 period H
 
                if self.saveSF:
                    alg.scaleFactorDecoration = f"muon_{self.prefixSF}_{trig_short}_%SYS%"
                if self.saveEff:
                    alg.mcEfficiencyDecoration = f"muon_{self.prefixEff}_{trig_short}_%SYS%"
                if self.saveEffData:
                    alg.dataEfficiencyDecoration = f"muon_{self.prefixEffData}_{trig_short}_%SYS%"
                alg.outOfValidity = 2 #silent
                alg.outOfValidityDeco = f"bad_eff_muontrig_{trig_short}"
                alg.muons = config.readName (self.containerName)
                alg.preselection = config.getPreselection (self.containerName, '')
                if self.saveSF:
                    config.addOutputVar (self.containerName, alg.scaleFactorDecoration, f"{self.prefixSF}_{trig_short}")
                if self.saveEff:
                    config.addOutputVar (self.containerName, alg.mcEfficiencyDecoration, f"{self.prefixEff}_{trig_short}")
                if self.saveEffData:
                    config.addOutputVar (self.containerName, alg.dataEfficiencyDecoration, f"{self.prefixEffData}_{trig_short}")
 
 
class MuonLRTMergedConfig (ConfigBlock) :
    def __init__ (self) :
        super (MuonLRTMergedConfig, self).__init__ ()
        self.addOption (
            'inputMuons', 'Muons', type=str,
            noneAction='error',
            info="the name of the input muon container."
        )
        self.addOption (
            'inputLRTMuons', 'MuonsLRT', type=str,
            noneAction='error',
            info="the name of the input LRT muon container."
        )
        self.addOption (
            'containerName', 'Muons_LRTMerged', type=str,
            noneAction='error',
            info="the name of the output container after LRT merging."
        )
    
    def instanceName (self) :
        return self.containerName
 
    def makeAlgs (self, config) :
 
        if config.isPhyslite() :
            raise(RuntimeError("Muon LRT merging is not available in Physlite mode"))
 
        alg = config.createAlgorithm( "CP::MuonLRTMergingAlg", "MuonLRTMergingAlg" )
        alg.PromptMuonLocation = self.inputMuons
        alg.LRTMuonLocation = self.inputLRTMuons
        alg.OutputMuonLocation = self.containerName
        alg.UseRun3WP = config.geometry() >= LHCPeriod.Run3 
        alg.CreateViewCollection = False
 
class MuonContainerMergingConfig (ConfigBlock) :
    def __init__ (self) :
        super (MuonContainerMergingConfig, self).__init__ ()
        self.addOption (
            'inputMuonContainers', [], type=list,
            noneAction='error',
            info="list of container names to be merged (of type `xAOD::MuonContainer`)."
        )
        self.addOption (
            'outputMuonLocation', 'MuonsMerged', type=str,
            noneAction='error',
            info="the name of the output muon container."
        )
        self.addOption (
            'createViewCollection', True, type=bool,
            info="whether the output container should be a view container rather than a deep copy."
        )
 
    def instanceName (self) :
        return self.outputMuonLocation
 
    def makeAlgs (self, config) :
        alg = config.createAlgorithm( "CP::MuonContainerMergingAlg", "MuonContainerMergingAlg" )
        alg.InputMuonContainers = self.inputMuonContainers
        alg.OutputMuonLocation = self.outputMuonLocation
        alg.CreateViewCollection = self.createViewCollection
 
@groupBlocks
def MuonWorkingPoint(seq):
    seq.append(MuonWorkingPointSelectionConfig())
    seq.append(MuonWorkingPointEfficiencyConfig())