TauCommonConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
 
def AddTauAugmentationCfg(flags, wp="GNTauVeryLoose", **kwargs):
    kwargs.setdefault("TauContainerName", "TauJets")
 
    acc = ComponentAccumulator()
 
    # tau selection relies on RNN electron veto, we must decorate the fixed eveto WPs before applying tau selection
    acc.merge(AddTauIDDecorationCfg(flags, TauContainerName=kwargs["TauContainerName"]))
 
    from DerivationFrameworkTools.DerivationFrameworkToolsConfig import AsgSelectionToolWrapperCfg
    from TauAnalysisTools.TauAnalysisToolsConfig import TauSelectionToolCfg
 
    TauAugmentationTools = []
 
    config = {
      "GNTauVeryLoose" : "TauAnalysisAlgorithms/tau_selection_gntau_veryloose_noeleid.conf",
      "GNTauLoose"     : "TauAnalysisAlgorithms/tau_selection_gntau_loose_noeleid.conf",
      "GNTauMedium"    : "TauAnalysisAlgorithms/tau_selection_gntau_medium_noeleid.conf",
      "GNTauTight"     : "TauAnalysisAlgorithms/tau_selection_gntau_tight_noeleid.conf",
    }
 
    TauSelector = acc.popToolsAndMerge(TauSelectionToolCfg(flags,
                                                           name = f'TauSelector{wp}',
                                                           ConfigPath = config[wp]))
    acc.addPublicTool(TauSelector)
 
    TauWrapper  = acc.getPrimaryAndMerge(AsgSelectionToolWrapperCfg(flags,
                                                                    name               = f"Tau{wp}Wrapper",
                                                                    AsgSelectionTool   = TauSelector,
                                                                    StoreGateEntryName = f"DFTau{wp}",
                                                                    ContainerName      = kwargs["TauContainerName"]))
    TauAugmentationTools.append(TauWrapper)
 
    CommonAugmentation = CompFactory.DerivationFramework.CommonAugmentation
    acc.addEventAlgo(CommonAugmentation(f"Tau{wp}AugmentationKernel", AugmentationTools = TauAugmentationTools))
    return acc
 
def AddTauIDDecorationCfg(flags, **kwargs):
    """Decorate tau ID scores and working points"""
 
    #kwargs.setdefault("evetoFix",         True)
    #kwargs.setdefault("GNNTauID",         True)
    tauContainerKey = kwargs.setdefault("TauContainerName", "TauJets")
    #kwargs.setdefault("prefix",           kwargs['TauContainerName'])
 
    acc = ComponentAccumulator()
 
    import tauRec.TauToolHolder as tauTools
    tools = []
    doEvetoWP = False
    scoreNames = []
    WPNames = []
 
    #def cacheToolProperties(tool):
    #    doEvetoWP = tool.UseAbsEta
    #    scoreName = [tool.ScoreName] if tools[-1].ScoreName != "RNNEleScore" else []
    #    wpName = tool.DecorWPNames
 
    if kwargs.pop('evetoFix', True):
        tools.append( acc.popToolsAndMerge(tauTools.TauWPDecoratorEleRNNCfg(flags)) )
        # Cache tool properties
        doEvetoWP |= tools[-1].UseAbsEta
        # The original RNNEleScore should not be overriden
        if tools[-1].ScoreName != "RNNEleScore": scoreNames.append(tools[-1].ScoreName)
        scoreNames.append(tools[-1].NewScoreName)
        WPNames += tools[-1].DecorWPNames
 
    if kwargs.pop('GNNTauID', True):
        # vertex-corrected clusters must be rebuilt for tau ID
        tools.append( acc.popToolsAndMerge(tauTools.TauVertexedClusterDecoratorCfg(flags)) )
        # Add in GNTau!
        # evaluate GNTau score for v0prune model
        tools.append( acc.popToolsAndMerge(tauTools.TauGNNEvaluatorCfg(flags,0,applyLooseTrackSel=True)) )
        # evaluate GNTau score for v1trunc model
        tools.append( acc.popToolsAndMerge(tauTools.TauGNNEvaluatorCfg(flags,1,applyLooseTrackSel=True)) )
        # set WPs decision for v0prune model
        tools.append( acc.popToolsAndMerge(tauTools.TauWPDecoratorGNNCfg(flags,0)) )
        # Cache tool properties
        doEvetoWP |= tools[-1].UseAbsEta
        # The original RNNEleScore should not be overriden
        if tools[-1].ScoreName != "RNNEleScore": scoreNames.append(tools[-1].ScoreName)
        scoreNames.append(tools[-1].NewScoreName)
        WPNames += tools[-1].DecorWPNames
 
        # set WPs decision for v1trunc model
        tools.append( acc.popToolsAndMerge(tauTools.TauWPDecoratorGNNCfg(flags,1)) )
        # Cache tool properties
        doEvetoWP |= tools[-1].UseAbsEta
        # The original RNNEleScore should not be overriden
        if tools[-1].ScoreName != "RNNEleScore": scoreNames.append(tools[-1].ScoreName)
        scoreNames.append(tools[-1].NewScoreName)
        WPNames += tools[-1].DecorWPNames
 
    if tools:
        kwargs.setdefault("DoEvetoWP", doEvetoWP)
        kwargs.setdefault("ScoreDecorationKeys", scoreNames)
        kwargs.setdefault("WPDecorationKeys", WPNames)
 
        for tool in tools:
            acc.addPublicTool(tool)
        kwargs.setdefault("TauIDTools", tools)
 
        TauIDDecoratorWrapper = CompFactory.DerivationFramework.TauIDDecoratorWrapper
        TauIDDecoratorKernel = CompFactory.DerivationFramework.CommonAugmentation
 
        prefix = kwargs.pop('prefix', tauContainerKey)
        tauIDDecoratorWrapper = TauIDDecoratorWrapper(name = f"{prefix}_TauIDDecoratorWrapper",
                                                      **kwargs)
        print("PXQW TauIDDecoratorsWrapper: " + str(tauIDDecoratorWrapper))
        acc.addPublicTool(tauIDDecoratorWrapper)
        acc.addEventAlgo(TauIDDecoratorKernel(name = f"{prefix}_TauIDDecorKernel",
                                              AugmentationTools = [tauIDDecoratorWrapper]))
 
    return acc
 
# Attach displaced Tau ID scores
def AddTauIDDisplacedDecorationCfg(flags, **kwargs):
    """Decorate displaced tau ID scores and working points. Follows AddTauIDDisplacedDecorationCfg()"""
    tauContainerKey = kwargs.setdefault("TauContainerName", "TauJetsLRT")
 
    flags_TauLRT = flags.cloneAndReplace("Tau.ActiveConfig", "Tau.TauLRT")
 
    acc = ComponentAccumulator()
 
    scoreNames = []
 
    import tauRec.TauToolHolder as tauTools
    tool_prompt = acc.popToolsAndMerge(tauTools.TauGNNEvaluatorCfg(flags,0,applyLooseTrackSel=True))
    scoreNames += ["GNTauScore_v0prune"]
 
    tool_displaced =  acc.popToolsAndMerge(tauTools.TauGNNDisplacedEvaluatorCfg(flags_TauLRT, tauContainerName=tauContainerKey))
    scoreNames += ["GNdTauScore", "GNdTauProbTau",  "GNdTauProbJet"]
 
    kwargs.setdefault("ScoreDecorationKeys", scoreNames)
    kwargs.setdefault("WPDecorationKeys", [])
 
    acc.addPublicTool(tool_prompt)
    acc.addPublicTool(tool_displaced)
    kwargs.setdefault("TauIDTools", [tool_prompt, tool_displaced])
 
    tauIDDecoratorWrapper = CompFactory.DerivationFramework.TauIDDecoratorWrapper(
        name = f"{tauContainerKey}_TauIDDisplacedDecoratorWrapper",
        **kwargs,
    )
    acc.addPublicTool(tauIDDecoratorWrapper)
 
    prefix = kwargs.pop('prefix', tauContainerKey)
    acc.addEventAlgo(
        CompFactory.DerivationFramework.CommonAugmentation(
            name = f"{prefix}_TauDisplacedIDDecorKernel",
            AugmentationTools = [tauIDDecoratorWrapper],
        )
    )
    return acc
 
# TauJets_MuonRM steering
def AddMuonRemovalTauAODReRecoAlgCfg(flags, **kwargs):
    """Configure the MuonRM AOD tau building"""
 
    acc = ComponentAccumulator()
    inputTauJets = kwargs.setdefault("Key_tauContainer", "TauJets")
    kwargs.setdefault("Key_tauOutputContainer", "TauJets_MuonRM")
    kwargs.setdefault("Key_pi0OutputContainer", "TauFinalPi0s_MuonRM")
    kwargs.setdefault("Key_neutralPFOOutputContainer", "TauNeutralParticleFlowObjects_MuonRM")
    kwargs.setdefault("Key_chargedPFOOutputContainer", "TauChargedParticleFlowObjects_MuonRM")
    kwargs.setdefault("Key_hadronicPFOOutputContainer", "TauHadronicParticleFlowObjects_MuonRM")
    kwargs.setdefault("Key_tauTrackOutputContainer", "TauTracks_MuonRM")
    kwargs.setdefault("Key_vertexOutputContainer", "TauSecondaryVertices_MuonRM")
 
    # get tools from holder
    import tauRec.TauToolHolder as tauTools
    if "modificationTools" not in kwargs:
        tools_mod = []
        tools_mod.append( acc.popToolsAndMerge(tauTools.TauAODMuonRemovalCfg(flags)) )
        for tool in tools_mod:
            tool.inAOD = True
        kwargs.setdefault("modificationTools", tools_mod)
 
    if "officialTools" not in kwargs:
        tools_after = []
        tools_after.append( acc.popToolsAndMerge(tauTools.TauVertexedClusterDecoratorCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauTrackRNNClassifierCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.EnergyCalibrationLCCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauCommonCalcVarsCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauSubstructureCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.Pi0ClusterCreatorCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.Pi0ClusterScalerCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.Pi0ScoreCalculatorCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.Pi0SelectorCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauVertexVariablesCfg(flags)) )
        import PanTauAlgs.JobOptions_Main_PanTau as pantau
        tools_after.append( acc.popToolsAndMerge(pantau.PanTauCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauCombinedTESCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.MvaTESVariableDecoratorCfg(flags)) )
        tools_after[-1].EventShapeKey = ''
        tools_after.append( acc.popToolsAndMerge(tauTools.MvaTESEvaluatorCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauIDVarCalculatorCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauJetRNNEvaluatorCfg(flags,applyLooseTrackSel=True)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauWPDecoratorJetRNNCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauEleRNNEvaluatorCfg(flags,applyLooseTrackSel=True )) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauWPDecoratorEleRNNCfg(flags)) )
        tools_after.append( acc.popToolsAndMerge(tauTools.TauDecayModeNNClassifierCfg(flags)) )
        for tool in tools_after:
            tool.inAOD = True
        kwargs.setdefault("officialTools", tools_after)
 
    kwargs.setdefault("ExtraInputs",
                      [ ( 'xAOD::TauJetContainer' , "StoreGateSvc+{baseName}.truthJetLink".format(baseName = inputTauJets)),
                        ( 'xAOD::TauJetContainer' , "StoreGateSvc+{baseName}.truthParticleLink".format(baseName = inputTauJets)) ] )
    TauAODRunnerAlg=CompFactory.getComp("TauAODRunnerAlg")
    myTauAODRunnerAlg = TauAODRunnerAlg(
        name = "MuonRemovalTauAODReRecoAlg",
        **kwargs
    )
    acc.addEventAlgo(myTauAODRunnerAlg)
    return acc
 
 
def TauThinningCfg(flags, name, **kwargs):
    """configure tau thinning"""
 
    acc = ComponentAccumulator()
    TauThinningTool = CompFactory.DerivationFramework.TauThinningTool
    acc.addPublicTool(TauThinningTool(name, **kwargs), primary=True)
    return acc