MuonPhaseII/MuonPatternRecognition/MuonTruthAlgsR4/python/MuonTruthAlgsConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
from AthenaConfiguration.AccumulatorCache import AccumulatorCache
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
 
 
def TruthSegmentMakerCfg(flags, name = "MuonTruthSegmentMaker", useSDO = True, **kwargs):
    result = ComponentAccumulator()
    if not flags.Input.isMC:
        return result
 
    from ActsAlignmentAlgs.AlignmentAlgsConfig import ActsGeometryContextAlgCfg
    result.merge(ActsGeometryContextAlgCfg(flags))
    containerNames = []
    
    if flags.Detector.EnableMDT: containerNames+=["MDT_SDO" if useSDO else "xMdtSimHits"]       
    if flags.Detector.EnableRPC: containerNames+=["RPC_SDO" if useSDO else "xRpcSimHits"]
    if flags.Detector.EnableTGC: containerNames+=["TGC_SDO" if useSDO else "xTgcSimHits"]
    if flags.Detector.EnableMM: containerNames+=["MM_SDO" if useSDO else "xMmSimHits"]
    if flags.Detector.EnablesTGC: containerNames+=["sTGC_SDO" if useSDO else "xStgcSimHits"] 
    
    from MuonConfig.MuonCalibrationConfig import MdtCalibDbAlgCfg, NswErrorCalibDbAlgCfg
    if flags.Detector.EnableMDT:
        result.merge(MdtCalibDbAlgCfg(flags))
    if flags.Detector.EnableMM or flags.Detector.EnablesTGC:
        result.merge(NswErrorCalibDbAlgCfg(flags))
    kwargs.setdefault("SimHitKeys", containerNames)
    kwargs.setdefault("includePileUpHits", flags.Muon.includePileUpTruth)
 
    the_alg = CompFactory.MuonR4.TruthSegmentMaker(name, **kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
def MeasToSimHitAssocAlgCfg(flags, name="MeasToSimHitConvAlg", **kwargs):
    result = ComponentAccumulator()
    from ActsAlignmentAlgs.AlignmentAlgsConfig import ActsGeometryContextAlgCfg
    result.merge(ActsGeometryContextAlgCfg(flags))
    the_alg = CompFactory.MuonR4.PrepDataToSimHitAssocAlg(name, **kwargs)
    result.addEventAlgo(the_alg, primary = True)    
    return result
 
def TruthHitAssociationCfg(flags, suffix = ""):
    result = ComponentAccumulator()
    if not flags.Input.isMC:
        return result
    from MuonConfig.MuonDataPrepConfig import PrimaryMeasContNamesCfg
    for cont_name in PrimaryMeasContNamesCfg(flags):
        simHits = ""
        if "xMdt" in cont_name: 
            simHits = "MDT_SDO"
        elif "xRpc" in cont_name:
            simHits = "RPC_SDO"
        elif "xTgc" in cont_name:
            simHits = "TGC_SDO"
        elif "MM" in cont_name:
            simHits = "MM_SDO"
        else:
            simHits = "sTGC_SDO"
        result.merge(MeasToSimHitAssocAlgCfg(flags,
                                             name=f"Muon{cont_name}ToSimHitAssoc{suffix}",
                                             SimHits = simHits,
                                             Measurements=cont_name))
        result.addEventAlgo(CompFactory.MuonR4.MeasTruthAssocMapProvider(name=f"Muon{cont_name}TruthAssocMap{suffix}",
                                                                         PrdContainer=cont_name,
                                                                         AssociationMap=f"{cont_name}_TruthAssocMap"))
    return result
 
def SimHitToTruthPartAlgCfg(flags, useSDO = False):
    result = ComponentAccumulator()
    if flags.Detector.GeometryMDT: 
        container = "xMdtSimHits" if not useSDO else "MDT_SDO"
        the_alg = CompFactory.MuonR4.SimHitToTruthPartAssocAlg("MuonMdtSimHitToTruthParticleAssoc",
                                                        SimHitContainer = container, 
                                                        HitIdDecoration = "truthMdtHits")
        result.addEventAlgo(the_alg)
    if flags.Detector.GeometryRPC:
        container = "xRpcSimHits" if not useSDO else "RPC_SDO"
        the_alg = CompFactory.MuonR4.SimHitToTruthPartAssocAlg("MuonRpcSimHitToTruthParticleAssoc",
                                                        SimHitContainer = container, 
                                                        HitIdDecoration = "truthRpcHits")
        result.addEventAlgo(the_alg)
    if flags.Detector.GeometryTGC: 
        container = "xTgcSimHits" if not useSDO else "TGC_SDO"
        the_alg = CompFactory.MuonR4.SimHitToTruthPartAssocAlg("MuonTgcSimHitToTruthParticleAssoc",
                                                        SimHitContainer = container,
                                                        HitIdDecoration = "truthTgcHits")
        result.addEventAlgo(the_alg)       
    if flags.Detector.GeometryMM: 
        container = "xMmSimHits" if not useSDO else "MM_SDO"
        the_alg = CompFactory.MuonR4.SimHitToTruthPartAssocAlg("MuonMmSimHitToTruthParticleAssoc",
                                                        SimHitContainer = container, 
                                                        HitIdDecoration = "truthMMHits")
        result.addEventAlgo(the_alg) 
    if flags.Detector.GeometrysTGC: 
        container = "xStgcSimHits" if not useSDO else "sTGC_SDO"
        the_alg = CompFactory.MuonR4.SimHitToTruthPartAssocAlg("MuonsTGCSimHitToTruthParticleAssoc",
                                                        SimHitContainer = container, 
                                                        HitIdDecoration = "truthStgcHits")
        result.addEventAlgo(the_alg) 
 
    return result
 
def TruthHitSummaryAlgCfg(flags, name="MuonTruthHitSummaryAlg", **kwargs):
    result = ComponentAccumulator()
    from MuonTrackFindingAlgs.TrackFindingConfig import TrackSummaryToolCfg, TrackSummaryLockCfg
    kwargs.setdefault("SummaryTool", result.popToolsAndMerge(TrackSummaryToolCfg(flags,  
                                                                fillHoles=False, fillOutliers=False, 
                                                                recomputeSegment =False)))
    
    kwargs.setdefault("ExtraOutputs", [('xAOD::TruthParticleContainer', 'MuonTruthParticles.truthHitSummary')])
    result.merge(TrackSummaryLockCfg(flags, inContainer="MuonTruthParticles", fillHoles = False, 
                                            fillOutliers = False, ExtraInputs = kwargs["ExtraOutputs"]))
    the_alg = CompFactory.MuonR4.TruthHitSummaryAlg(name, **kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
def RecoSegmentTruthAssocCfg(flags, **kwargs):
    result = ComponentAccumulator()
    
    kwargs.setdefault("SegmentKey", "")
    from MuonConfig.MuonDataPrepConfig import PrimaryMeasContNamesCfg
    kwargs.setdefault("PrdContainer", PrimaryMeasContNamesCfg(flags))
    the_alg = CompFactory.MuonR4.RecoSegToTruthAssocAlg(**kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
def TruthSegmentToTruthPartAssocCfg(flags, name="MuonTruthSegmentToTruthAssocAlg", **kwargs):
    result = ComponentAccumulator()
    hitDecors = []
    if flags.Detector.GeometryMDT: hitDecors+=["truthMdtHits"]
    if flags.Detector.GeometryRPC: hitDecors+=["truthRpcHits"]
    if flags.Detector.GeometryTGC: hitDecors+=["truthTgcHits"]
    if flags.Detector.GeometryMM: hitDecors+=["truthMMHits"]
    if flags.Detector.GeometrysTGC: hitDecors+=["truthStgcHits"]
    kwargs.setdefault("SimHitIds", hitDecors)
    from ActsConfig.ActsGeometryConfig import ActsExtrapolationToolCfg
    kwargs.setdefault("ExtrapolationTool", result.popToolsAndMerge(ActsExtrapolationToolCfg(flags, MaxSteps=100000)))
    kwargs.setdefault("includePileUpObjs", flags.Muon.includePileUpTruth)
    the_alg = CompFactory.MuonR4.TruthSegToTruthPartAssocAlg(name, **kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
def TrackToTruthPartAssocCfg(flags, **kwargs):
    result = ComponentAccumulator()
    hitDecors = []
    if flags.Detector.GeometryMDT: hitDecors+=["truthMdtHits"]
    if flags.Detector.GeometryRPC: hitDecors+=["truthRpcHits"]
    if flags.Detector.GeometryTGC: hitDecors+=["truthTgcHits"]
    if flags.Detector.GeometryMM: hitDecors+=["truthMMHits"]
    if flags.Detector.GeometrysTGC: hitDecors+=["truthStgcHits"]
    kwargs.setdefault("SimHitIds", hitDecors)
    the_alg = CompFactory.MuonR4.TrackToTruthPartAssocAlg(**kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
 
def MuonTruthSegConnectorAlgCfg(flags, name="MuonTruthSegConnector",**kwargs):
    result = ComponentAccumulator()
    the_alg = CompFactory.MuonR4.TruthSegConnectionAlg(name=name, **kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
 
def MuonToTruthAssocAlgCfg(flags, name="MuonToTruthMatchingAlg", **kwargs):
    result = ComponentAccumulator()
    if not flags.Muon.setupTruthAlgorithms:
        return result
    the_alg = CompFactory.MuonR4.MuonToTruthAssocAlg(name, **kwargs)
    result.addEventAlgo(the_alg, primary = True)
    return result
 
# Fragment for algs producing (and decorating) truth objects
def MuonTruthObjCreatorsCfg(flags, useSDO=True):
    result = ComponentAccumulator()
    if not flags.Muon.setupTruthAlgorithms:
        return result
 
    from MuonConfig.MuonTruthAlgsConfig import TruthMuonMakerAlgCfg
    result.merge(TruthMuonMakerAlgCfg(flags))
    result.merge(SimHitToTruthPartAlgCfg(flags, useSDO = useSDO))
    result.merge(TruthSegmentMakerCfg(flags, useSDO = useSDO))
    result.merge(TruthSegmentToTruthPartAssocCfg(flags))
    result.merge(TruthHitSummaryAlgCfg(flags))
 
    if flags.Muon.includePileUpTruth:
        result.merge(MuonTruthSegConnectorAlgCfg(flags))
 
    return result
 
@AccumulatorCache
def MuonTruthAlgsCfg(flags, useSDO=True, recoAssoc = True):
    result = ComponentAccumulator()
    if not flags.Muon.setupTruthAlgorithms:
        return result
    
    result.merge(MuonTruthObjCreatorsCfg(flags, useSDO=useSDO))
 
    if useSDO and recoAssoc:
        result.merge(TruthHitAssociationCfg(flags))
   
    return result
 
def MuonActsTruthMatchingCfg(flags, 
                             TrackContainer:str,
                             ParticleContainer:str,
                             isCombined: bool = False):
    result = ComponentAccumulator()
    from ActsConfig.ActsTruthConfig import ActsTrackParticleTruthDecorationAlgCfg, \
                                           ActsTrackToTruthAssociationAlgCfg, \
                                           ActsTruthParticleHitCountAlgCfg
    MapSuffix="MS" if not isCombined else "CB"                                     
    TruthHitCountMap = f"ActsHitCountMap_{MapSuffix}"
    
    result.merge(ActsTruthParticleHitCountAlgCfg(flags, name=f"MuonActsTruthHitsCounts_{MapSuffix}",
                                                 TruthParticleHitCountsOut= TruthHitCountMap,
                                                 doInner = isCombined, doMuon=True))
    
    result.merge(ActsTrackToTruthAssociationAlgCfg(flags, name =f"MuonTrackTruthMatching_{TrackContainer}",
                                                    ACTSTracksLocation=TrackContainer,
                                                    AssociationMapOut=f"{TrackContainer}ToTruthParticles",
                                                    doInner = isCombined, doMuon = True))
    
    result.merge(ActsTrackParticleTruthDecorationAlgCfg(flags, name=f"MuonTrackParticleToTruth_{ParticleContainer}",
                                                        TrackParticleContainerName=ParticleContainer,
                                                        TrackToTruthAssociationMaps =[f"{TrackContainer}ToTruthParticles"],
                                                        DerivateTruthKey="MuonTruthParticles",
                                                        TruthParticleHitCounts=TruthHitCountMap))
    return result