DerivationTrackIsoConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
 
__doc__ = """
          Instantiate more custom track isolation in derivation
          """
 
from AthenaCommon.Logging import logging
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
 
def DerivationTrackIsoCfg(flags,**jwarg):
 
    mlog = logging.getLogger('DerivationTrackIsoConfig')
    mlog.info('Starting track isolation for derivation')
 
    acc = ComponentAccumulator()
 
    from xAODPrimitives.xAODIso import xAODIso as isoPar
    # dR in decreasing order
    ptcone_list = [ [ isoPar.ptcone40, isoPar.ptcone30, isoPar.ptcone20 ] ]
    trkcor_list = [ [ isoPar.coreTrackPtr ] ]
 
    if 'object_types' not in jwarg:
        object_types = ("Electrons", "Photons", "Muons")
    else:
        object_types = jwarg['object_types']
 
    if 'WP' not in jwarg:
        WP = 'Nonprompt_All_MaxWeight'
    else:
        WP = jwarg['WP']
        # TrackIsolationTool need a specific track-vertex association tool
        #   if WP is not Nonprompt_All_MaxWeight
        from TrackVertexAssociationTool.TrackVertexAssociationToolConfig import isoTTVAToolCfg
        ttvaCA = acc.popToolsAndMerge(
            isoTTVAToolCfg(flags, WorkingPoint = WP))
 
    postfix = ''
    if 'postfix' in jwarg:
        postfix = jwarg['postfix']
 
    if 'ptCuts' not in jwarg:
        ptCuts = (500, 1000)
    else:
        ptCuts = jwarg['ptCuts']
 
    do_egamma = any(x in object_types for x in ("Electrons", "Photons"))
 
    mlog.info('will isolate objects '+" ".join(object_types)+' with TTVA WP = '+WP)
 
    for track_pt in ptCuts:
        for loose_cone in (True, False):
            if loose_cone and not do_egamma:
                # Loose cone isolation variables only for electrons and photons
                continue
 
            cone_str = "LooseCone" if loose_cone else ""
 
            # The custom variable suffix
            name = f"{WP}TTVA{cone_str}_pt{track_pt}"
 
            # Build up extra IsolationBuilder kwargs
            kwargs = {}
            if "Electrons" in object_types:
                kwargs["ElIsoTypes"] = ptcone_list
                kwargs["ElCorTypes"] = trkcor_list
                kwargs["ElCorTypesExtra"] = [[]]
                kwargs["CustomConfigurationNameEl"] = name
            if "Electrons" in object_types and "LRT" in postfix:
                kwargs["ElectronCollectionContainerName"] = "LRTElectrons"
            if "Photons" in object_types:
                kwargs["PhIsoTypes"] = ptcone_list
                kwargs["PhCorTypes"] = trkcor_list
                kwargs["PhCorTypesExtra"] = [[]]
                kwargs["CustomConfigurationNamePh"] = name
            if "Muons" in object_types and not loose_cone:
                kwargs["MuIsoTypes"] = ptcone_list
                kwargs["MuCorTypes"] = trkcor_list
                kwargs["MuCorTypesExtra"] = [[]]
                kwargs["CustomConfigurationNameMu"] = name
            if "Muons" in object_types and not loose_cone and "LRT" in postfix:
                kwargs["MuonCollectionContainerName"] = "MuonsLRT"
                
            
            toolkwargs = {}
            # the TTVA tool is not the default choice
            if 'WP' in jwarg:
                toolkwargs['TTVATool'] = ttvaCA
            # and a track selection tool
 
            from InDetConfig.InDetTrackSelectionToolConfig import (
                isoTrackSelectionToolCfg )
            toolkwargs['TrackSelectionTool'] = acc.popToolsAndMerge(
                isoTrackSelectionToolCfg(flags, minPt = track_pt))
 
            if loose_cone:
                toolkwargs['CoreTrackEtaRange'] = 0.01
 
            from IsolationAlgs.IsoToolsConfig import TrackIsolationToolCfg
            kwargs['TrackIsolationTool'] = acc.popToolsAndMerge(
                TrackIsolationToolCfg(flags,**toolkwargs))
 
            # The algorithm
            kwargs['name'] = f"IsolationBuilder{WP}{cone_str}{track_pt}{postfix}"
            acc.addEventAlgo(CompFactory.IsolationBuilder(**kwargs))
 
    return acc
 
def iso_vars(WP = "Nonprompt_All_MaxWeight"):
    # Get the list of isolation variables calculated by these functions
    iso_vars = []
    for track_pt in 500, 1000:
        for cone_str in "", "LooseCone":
            name = f"{WP}TTVA{cone_str}_pt{track_pt}"
            iso_vars += [ "ptconeCorrBitset_"+name,
                             "ptconecoreTrackPtrCorrection_"+name ]
            for cone_size in 20, 30, 40:
                for var_str in "", "var":
                    iso_vars.append(f"pt{var_str}cone{cone_size}_{name}")
    return iso_vars
 
if __name__ == "__main__":
    from AthenaConfiguration.AllConfigFlags import initConfigFlags
    from AthenaConfiguration.ComponentAccumulator import printProperties
    from AthenaConfiguration.TestDefaults import defaultTestFiles
    from AthenaConfiguration.MainServicesConfig import MainServicesCfg
    flags = initConfigFlags()
    flags.Input.Files = defaultTestFiles.RDO_RUN2
    flags.Output.doWriteESD = True  # To test the ESD parts
    flags.Output.doWriteAOD = True  # To test the AOD parts
    flags.lock()
 
    mlog = logging.getLogger("derivationTrackIsolationConfigTest")
    mlog.info("Configuring derivation track isolation: ")
 
    acc = MainServicesCfg(flags)
    acc.merge(DerivationTrackIsoCfg(flags))
    # print config of some of the algs :
    for par in ('500', '1000', 'LooseCone1000'):
        print('------ Config for ',par,' ----------')
        name = "IsolationBuilderNonprompt_All_MaxWeight"+par
        printProperties(mlog,
                        acc.getEventAlgo(name),
                        nestLevel=1,
                        printDefaults=True)