ActsTrackFittingConfig.py

Go to the documentation of this file.

#  Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
 
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
from ActsConfig.ActsConfigFlags import TrackFitterType
from ActsInterop import UnitConstants
from typing import Optional
 
def ActsToTrkFitterCfg(flags,
                       name: str = "ActsToTrkFitterTool",
                       fitterKind:  Optional[TrackFitterType] = None, 
                       **kwargs) -> ComponentAccumulator:
    
    acc = ComponentAccumulator()
 
    fitterKind = flags.Acts.trackFitterType  if fitterKind is None else fitterKind
 
    from ActsAlignmentAlgs.AlignmentAlgsConfig import ActsGeometryContextAlgCfg
    from ActsConfig.ActsGeometryConfig import ActsTrackingGeometrySvcCfg, ActsGeometryRealmConvToolCfg
    acc.merge(ActsGeometryContextAlgCfg(flags))
    acc.merge(ActsTrackingGeometrySvcCfg(flags))
    kwargs.setdefault("GeometryRealmConvTool", acc.getPrimaryAndMerge(ActsGeometryRealmConvToolCfg(flags)))
    if "ATLASConverterTool" not in kwargs:
        from ActsConfig.ActsEventCnvConfig import ActsToTrkConverterToolCfg
        kwargs.setdefault('ATLASConverterTool', acc.getPrimaryAndMerge(ActsToTrkConverterToolCfg(flags)))
 
    if fitterKind is TrackFitterType.KalmanFitter:
        kwargs.setdefault('SeedCovarianceScale', 100.0)
    
    if "ActsFitterTool" not in kwargs:
        if fitterKind is TrackFitterType.KalmanFitter:
            kwargs.setdefault("ActsFitterTool", acc.popToolsAndMerge(ActsFitterCfg(flags, 
                                                                                   fitterKind=fitterKind, 
                                                                                   UseDirectNavigation = False)))
        else:
            kwargs.setdefault("ActsFitterTool", acc.popToolsAndMerge(ActsFitterCfg(flags, 
                                                                                   fitterKind=fitterKind)))
 
    acc.setPrivateTools(CompFactory.ActsTrk.ActsToTrkFitterWrapTool(name, **kwargs))
 
    return acc
 
def ActsFitterCfg(flags,
                  name: str = "ActsKalmanFitterTool",
                  fitterKind:  Optional[TrackFitterType] = None, 
                  **kwargs) -> ComponentAccumulator:
    acc = ComponentAccumulator()
 
    fitterKind = flags.Acts.trackFitterType  if fitterKind is None else fitterKind
 
    if flags.Acts.fitFromPRD:       
        if flags.Detector.GeometryITk:
            from TrkConfig.TrkRIO_OnTrackCreatorConfig import ITkRotCreatorCfg
            kwargs.setdefault('RotCreatorTool', acc.popToolsAndMerge(ITkRotCreatorCfg(flags)))
        elif flags.Detector.GeometryID:
            from TrkConfig.TrkRIO_OnTrackCreatorConfig import InDetRotCreatorCfg
            kwargs.setdefault('RotCreatorTool', acc.popToolsAndMerge(InDetRotCreatorCfg(flags)))
 
 
    # Make sure this is set correctly!
    #  /eos/project-a/acts/public/MaterialMaps/ATLAS/material-maps-Pixel-SCT.json
    from ActsAlignmentAlgs.AlignmentAlgsConfig import ActsGeometryContextAlgCfg
    from ActsConfig.ActsGeometryConfig import ActsTrackingGeometrySvcCfg, ActsGeometryRealmConvToolCfg
    acc.merge(ActsGeometryContextAlgCfg(flags))
    acc.merge(ActsTrackingGeometrySvcCfg(flags))
    kwargs.setdefault("GeometryRealmConvTool", acc.getPrimaryAndMerge(ActsGeometryRealmConvToolCfg(flags)))
 
    if fitterKind is TrackFitterType.KalmanFitter:
        kwargs.setdefault("ReverseFilteringPt", 1.0 * UnitConstants.GeV)
 
    if fitterKind is TrackFitterType.KalmanFitter:    # This flag is by default set to KalmanFitter
        acc.setPrivateTools(CompFactory.ActsTrk.KalmanFitterTool(name, **kwargs))
    elif fitterKind is TrackFitterType.GaussianSumFitter:
        name = name.replace("KalmanFitter", "GaussianSumFitter")
        acc.setPrivateTools(CompFactory.ActsTrk.GaussianSumFitterTool(name, **kwargs))
    elif fitterKind is TrackFitterType.GlobalChiSquareFitter:
        name = name.replace("KalmanFitter", "GlobalChiSquareFitter")
        acc.setPrivateTools(CompFactory.ActsTrk.GlobalChiSquareFitterTool(name, **kwargs))
 
    return acc
 
 
def ActsReFitterAlgCfg(flags,
                       name : str = "ActsReFitterAlg",
                       **kwargs) -> ComponentAccumulator:
    acc = ComponentAccumulator()
    kwargs.setdefault("ActsFitter", acc.popToolsAndMerge(ActsToTrkFitterCfg(flags)))
    kwargs.setdefault("TrackName", "ResolvedTracks")
    kwargs.setdefault("NewTrackName", "Refitted_Tracks")
    kwargs.setdefault("DoReFitFromPRD", flags.Acts.fitFromPRD)
    acc.addEventAlgo(CompFactory.ActsTrk.ReFitterAlg(name, **kwargs))
 
    if flags.Acts.writeTrackCollection:
        acc.merge(writeAdditionalTracks(flags))
 
    return acc
 
def forceITkActsReFitterAlgCfg(flags) -> ComponentAccumulator:
    #Use this flag in the --postInclude of ActsKfRefiting.sh to fit from the PRD (uncalibrated); Else to fit from the ROT (calibrated), use `(...).ActsReFitterAlgCfg` flag directly
   flags = flags.cloneAndReplace("Tracking.ActiveConfig",
                                 flags.Tracking.ITkPrimaryPassConfig.value)
   return ActsReFitterAlgCfg(flags)
 
 
def writeAdditionalTracks(flags,
                          trackName: str = 'ResolvedTracks',
                          newTrackName: str ='ReFitted_Tracks') -> ComponentAccumulator:
    acc = ComponentAccumulator()
    
    from xAODTrackingCnv.xAODTrackingCnvConfig import ITkTrackParticleCnvAlgCfg
 
    if flags.Tracking.doTruth:
        from InDetConfig.ITkTrackTruthConfig import ITkTrackTruthCfg
        acc.merge(ITkTrackTruthCfg(flags,
                                   Tracks = trackName,
                                   DetailedTruth = f"{trackName}DetailedTruth",
                                   TracksTruth = f"{trackName}TruthCollection"))
        acc.merge(ITkTrackTruthCfg(flags,
                                   Tracks = newTrackName,
                                   DetailedTruth = f"{newTrackName}DetailedTruth",
                                   TracksTruth = f"{newTrackName}TruthCollection"))
 
    acc.merge(ITkTrackParticleCnvAlgCfg(flags,
                                        name = f"{trackName}TrackParticleCnvAlg",
                                        TrackContainerName = trackName,
                                        xAODTrackParticlesFromTracksContainerName = f"{trackName}TrackParticles",
                                        TrackTruthContainerName = f"{trackName}TruthCollection")) 
    acc.merge(ITkTrackParticleCnvAlgCfg(flags,
                                        name = f"{newTrackName}TrackParticleCnvAlg",
                                        TrackContainerName = newTrackName,
                                        xAODTrackParticlesFromTracksContainerName = f"{newTrackName}TrackParticles",
                                        TrackTruthContainerName = f"{newTrackName}TruthCollection")) 
    
    from OutputStreamAthenaPool.OutputStreamConfig import addToESD, addToAOD
    itemList = [f"xAOD::TrackParticleContainer#{trackName}TrackParticles",
                f"xAOD::TrackParticleAuxContainer#{trackName}TrackParticlesAux.",
                f"xAOD::TrackParticleContainer#{newTrackName}TrackParticles",
                f"xAOD::TrackParticleAuxContainer#{newTrackName}TrackParticlesAux."]
 
    acc.merge(addToESD(flags, itemList))
    acc.merge(addToAOD(flags, itemList))
    return acc