IsolationBuilderConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
 
""" Instantiate the electron or photon isolation """
 
from AthenaCommon.Logging import logging
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
 
def egIsolationCfg(flags, name='egIsolation', noCalo=False, **kwargs):
 
    mlog = logging.getLogger(name)
    mlog.info('Starting '+name+' configuration')
 
    acc = ComponentAccumulator()
 
    suff = ''
    if name.find('photon') >= 0:
        suff = 'photon'
        pref = 'Ph'
    elif name.find('electron') >= 0:
        suff = 'electron'
        pref = 'El'
    else:
        mlog.error('Name '+name+' should contain electron or photon')
        return acc
        
    from xAODPrimitives.xAODIso import xAODIso as isoPar
    from IsolationAlgs.IsoToolsConfig import (
        TrackIsolationToolCfg, ElectronTrackIsolationToolCfg,
        EGammaCaloIsolationToolCfg)
 
    isoType  = []
    isoCor   = []
    isoExCor = []
 
    if flags.Detector.EnableID or flags.Detector.EnableITk:
        isoType.append([ isoPar.ptcone30, isoPar.ptcone20 ])
        isoCor.append([ isoPar.coreTrackPtr ])
        isoExCor.append([])
        if 'TrackIsolationTool' not in kwargs:
            if pref == 'Ph':
                kwargs['TrackIsolationTool'] = acc.popToolsAndMerge(
                    TrackIsolationToolCfg(flags))
            else:
                kwargs['TrackIsolationTool'] = acc.popToolsAndMerge(
                    ElectronTrackIsolationToolCfg(flags))
            
    if flags.Detector.EnableCalo and not noCalo:
        isoType.append(
            [ isoPar.topoetcone20, isoPar.topoetcone30, isoPar.topoetcone40 ])
        isoCor.append(
            [ isoPar.core57cells, isoPar.ptCorrection, isoPar.pileupCorrection ])
        # do not do pileup correction if HI with subtracted clusters
        if flags.HeavyIon.Egamma.doSubtractedClusters:
            isoCor[-1] = [ x for x in isoCor[-1] if x != isoPar.pileupCorrection ]
        isoExCor.append([])
        if 'CaloTopoIsolationTool' not in kwargs:
            kwargs['CaloTopoIsolationTool'] = acc.popToolsAndMerge(
                EGammaCaloIsolationToolCfg(flags))
 
    kwargs[f'{pref}IsoTypes'] = isoType
    kwargs[f'{pref}CorTypes'] = isoCor
    kwargs[f'{pref}CorTypesExtra'] = isoExCor
 
    kwargs['name'] = suff+'IsolationBuilder'
    
    acc.addEventAlgo(CompFactory.IsolationBuilder(**kwargs))
 
    mlog.info(suff+" isolation configured")
 
    return acc
 
def muIsolationCfg(flags, name='muIsolation', noCalo=False, **kwargs):
 
    mlog = logging.getLogger(name)
    mlog.info('Starting '+name+' configuration')
 
    acc = ComponentAccumulator()
 
    from xAODPrimitives.xAODIso import xAODIso as isoPar
    from IsolationAlgs.IsoToolsConfig import (
        TrackIsolationToolCfg, MuonCaloIsolationToolCfg)
 
    isoType  = []
    isoCor   = []
    isoExCor = []
 
    if flags.Detector.EnableID or flags.Detector.EnableITk:
        isoType.append([ isoPar.ptcone40, isoPar.ptcone30, isoPar.ptcone20 ])
        isoCor.append([ isoPar.coreTrackPtr ])
        isoExCor.append([])
        if 'TrackIsolationTool' not in kwargs:
            kwargs['TrackIsolationTool'] = acc.popToolsAndMerge(
                TrackIsolationToolCfg(flags))
        
    if flags.Detector.EnableCalo and not noCalo:
        isoType.append(
            [ isoPar.topoetcone20, isoPar.topoetcone30, isoPar.topoetcone40 ])
        isoCor.append([ isoPar.coreCone, isoPar.pileupCorrection ])
        # do not do pileup correction if HI with subtracted clusters
        if flags.HeavyIon.Egamma.doSubtractedClusters:
            isoCor[-1] = [ x for x in isoCor[-1] if x != isoPar.pileupCorrection ]
        isoExCor.append([])
        if ('CaloTopoIsolationTool' not in kwargs) or (
                'PFlowIsolationTool' not in kwargs):
            cisoTool = acc.popToolsAndMerge(MuonCaloIsolationToolCfg(flags))
        if 'CaloTopoIsolationTool' not in kwargs:
            kwargs['CaloTopoIsolationTool'] = cisoTool
        if flags.Reco.EnablePFlow and 'PFlowIsolationTool' not in kwargs:
            isoType.append(
                [ isoPar.neflowisol20, isoPar.neflowisol30, isoPar.neflowisol40 ])
            isoCor.append([ isoPar.coreCone, isoPar.pileupCorrection ])
            # do not do pileup correction for HI reco...
            # and pflow iso is more than experimental in such reco anyway
            if flags.HeavyIon.Egamma.doSubtractedClusters:
                isoCor[-1] = [ x for x in isoCor[-1] if x != isoPar.pileupCorrection ]
            isoExCor.append([])
            kwargs['PFlowIsolationTool'] = cisoTool
 
    kwargs['MuIsoTypes'] = isoType
    kwargs['MuCorTypes'] = isoCor
    kwargs['MuCorTypesExtra'] = isoExCor
    kwargs['name'] = 'muonIsolationBuilder'
    
    acc.addEventAlgo(CompFactory.IsolationBuilder(**kwargs))
 
    mlog.info("muon isolation configured")
 
    return acc
 
 
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("isolationConfigTest")
    mlog.info("Configuring photon isolation: ")
 
    acc = MainServicesCfg(flags)
    acc.merge(egIsolationCfg(flags,name = 'photonIsolation'))
    acc.merge(egIsolationCfg(flags,name = 'electronIsolation'))
    acc.merge(muIsolationCfg(flags,name = 'muonIsolation'))
    printProperties(mlog,
                    acc.getEventAlgo("photonIsolationBuilder"),
                    nestLevel=1,
                    printDefaults=True)
    printProperties(mlog,
                    acc.getEventAlgo("electronIsolationBuilder"),
                    nestLevel=1,
                    printDefaults=True)
    printProperties(mlog,
                    acc.getEventAlgo("muonIsolationBuilder"),
                    nestLevel=1,
                    printDefaults=True)
 
    with open("isolationconfig.pkl", "wb") as f:
        acc.store(f)