InDetGNNTrackingConfig.py

Go to the documentation of this file.

#
#  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
#
 
from pathlib import Path
 
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
 
 
def DumpObjectsCfg(
        flags, name="DumpObjects", outfile="Dump_GNN4Itk.root", **kwargs):
    '''
    create algorithm which dumps GNN training information to ROOT file
    '''
    acc = ComponentAccumulator()
 
    acc.addService(
        CompFactory.THistSvc(
            Output=[f"{name} DATAFILE='{outfile}', OPT='RECREATE'"]
        )
    )
 
    kwargs.setdefault("NtupleFileName", "/DumpObjects/")
    kwargs.setdefault("NtupleTreeName", "GNN4ITk")
    kwargs.setdefault("rootFile", True)
 
    acc.addEventAlgo(CompFactory.InDet.DumpObjects(name, **kwargs))
    return acc
 
def GNNTrackFinderToolCfg(flags, name='GNNTrackFinderTool', **kwargs):
    """Sets up a GNNTrackFinderTool tool and returns it."""
    acc = ComponentAccumulator()
    
    
    kwargs.setdefault("embeddingDim", 8)
    kwargs.setdefault("rVal", 1.7)
    kwargs.setdefault("knnVal", 500)
    kwargs.setdefault("filterCut", 0.21)
    kwargs.setdefault("inputMLModelDir", "TrainedMLModels4ITk")
    kwargs.setdefault("UseCUDA", False)
 
    from AthOnnxComps.OnnxRuntimeInferenceConfig import OnnxRuntimeInferenceToolCfg
    kwargs.setdefault("Embedding", acc.popToolsAndMerge(
        OnnxRuntimeInferenceToolCfg(flags, Path("TrainedMLModels4ITk") / "embedding.onnx")
    ))
    kwargs.setdefault("Filtering", acc.popToolsAndMerge(
        OnnxRuntimeInferenceToolCfg(flags, Path("TrainedMLModels4ITk") / "filtering.onnx")
    ))
    kwargs.setdefault("GNN", acc.popToolsAndMerge(
        OnnxRuntimeInferenceToolCfg(flags, Path("TrainedMLModels4ITk") / "gnn.onnx")
    ))
    
    acc.setPrivateTools(CompFactory.InDet.SiGNNTrackFinderTool(name, **kwargs))
    return acc
 
 
def SeedFitterToolCfg(flags, name="SeedFitterTool", **kwargs):
    """Sets up a SeedFitter tool and returns it."""
    acc = ComponentAccumulator()
    
    
    acc.setPrivateTools(CompFactory.InDet.SeedFitterTool(name, **kwargs))
    return acc
 
def GNNTrackReaderToolCfg(flags, name='GNNTrackReaderTool', **kwargs):
    """Set up a GNNTrackReader tool and return it."""
    acc = ComponentAccumulator()
 
    
    kwargs.setdefault("inputTracksDir", "gnntracks")
    kwargs.setdefault("csvPrefix", "track")
 
    acc.setPrivateTools(CompFactory.InDet.GNNTrackReaderTool(name, **kwargs))
    return acc
 
def GNNTrackMakerCfg(flags, name="GNNTrackMaker", **kwargs):
    """Sets up a GNNTrackMaker algorithm and returns it."""
    
    if flags.Tracking.GNN.usePixelHitsOnly:
        return GNNSeedingTrackMakerCfg(flags, name, **kwargs)
    else:
        return GNNEndToEndTrackMaker(flags, name, **kwargs)
 
def GNNEndToEndTrackMaker(flags, name="GNNEndToEndTrackMaker", **kwargs):
    """Sets up a GNNTrackMaker algorithm and returns it."""
    
    acc = ComponentAccumulator()
    
    
    SeedFitterTool = acc.popToolsAndMerge(SeedFitterToolCfg(flags))
    kwargs.setdefault("SeedFitterTool", SeedFitterTool)
    
    from TrkConfig.CommonTrackFitterConfig import ITkTrackFitterCfg
    InDetTrackFitter = acc.popToolsAndMerge(ITkTrackFitterCfg(flags))
    kwargs.setdefault("TrackFitter", InDetTrackFitter)
 
    if flags.Tracking.GNN.useTrackFinder:
        InDetGNNTrackFinderTool = acc.popToolsAndMerge(GNNTrackFinderToolCfg(flags))
        kwargs.setdefault("GNNTrackFinderTool", InDetGNNTrackFinderTool)
        kwargs.setdefault("GNNTrackReaderTool", None)
    elif flags.Tracking.GNN.useTrackReader:
        InDetGNNTrackReader = acc.popToolsAndMerge(GNNTrackReaderToolCfg(flags))
        kwargs.setdefault("GNNTrackReaderTool", InDetGNNTrackReader)
        kwargs.setdefault("GNNTrackFinderTool", None)
    else:
        raise RuntimeError("GNNTrackFinder or GNNTrackReader must be enabled!")
 
    acc.addEventAlgo(CompFactory.InDet.SiSPGNNTrackMaker(name, **kwargs))
    return acc
 
def GNNSeedingTrackMakerCfg(flags, name="GNNSeedingTrackMaker", **kwargs):
    """Sets up a GNN for seeding algorithm and returns it."""
    acc = ComponentAccumulator()
 
    from InDetConfig.SiCombinatorialTrackFinderToolConfig import SiDetElementBoundaryLinksCondAlg_xk_ITkPixel_Cfg, SiDetElementBoundaryLinksCondAlg_xk_ITkStrip_Cfg
    acc.merge(SiDetElementBoundaryLinksCondAlg_xk_ITkPixel_Cfg(flags))
    acc.merge(SiDetElementBoundaryLinksCondAlg_xk_ITkStrip_Cfg(flags))
 
    # To produce AtlasFieldCacheCondObj
    from MagFieldServices.MagFieldServicesConfig import (
        AtlasFieldCacheCondAlgCfg)
    acc.merge(AtlasFieldCacheCondAlgCfg(flags))
 
    from TrkConfig.TrkRIO_OnTrackCreatorConfig import ITkRotCreatorCfg
    ITkRotCreator = acc.popToolsAndMerge(ITkRotCreatorCfg(
        flags, name="ITkRotCreator"+flags.Tracking.ActiveConfig.extension))
    acc.addPublicTool(ITkRotCreator)
    kwargs.setdefault("RIOonTrackTool", ITkRotCreator)
 
    from TrkConfig.TrkExRungeKuttaPropagatorConfig import (
        RungeKuttaPropagatorCfg)
    ITkPatternPropagator = acc.popToolsAndMerge(
        RungeKuttaPropagatorCfg(flags, name="ITkPatternPropagator"))
    acc.addPublicTool(ITkPatternPropagator)
    kwargs.setdefault("PropagatorTool", ITkPatternPropagator)
 
    from TrkConfig.TrkMeasurementUpdatorConfig import KalmanUpdator_xkCfg
    ITkPatternUpdator = acc.popToolsAndMerge(
        KalmanUpdator_xkCfg(flags, name="ITkPatternUpdator"))
    acc.addPublicTool(ITkPatternUpdator)
    kwargs.setdefault("UpdatorTool", ITkPatternUpdator)
 
    from InDetConfig.InDetBoundaryCheckToolConfig import ITkBoundaryCheckToolCfg
    kwargs.setdefault("BoundaryCheckTool", acc.popToolsAndMerge(
        ITkBoundaryCheckToolCfg(flags)))
 
    from PixelConditionsTools.ITkPixelConditionsSummaryConfig import (
        ITkPixelConditionsSummaryCfg)
    kwargs.setdefault("PixelSummaryTool", acc.popToolsAndMerge(
        ITkPixelConditionsSummaryCfg(flags)))
 
    from SCT_ConditionsTools.ITkStripConditionsToolsConfig import (
        ITkStripConditionsSummaryToolCfg)
    kwargs.setdefault("StripSummaryTool", acc.popToolsAndMerge(
        ITkStripConditionsSummaryToolCfg(flags)))
 
    if flags.Tracking.GNN.useTrackFinder:
        kwargs.setdefault("GNNTrackFinderTool", acc.popToolsAndMerge(GNNTrackFinderToolCfg(flags)))
        kwargs.setdefault("GNNTrackReaderTool", None)
    elif flags.Tracking.GNN.useTrackReader:
        kwargs.setdefault("GNNTrackReaderTool", acc.popToolsAndMerge(GNNTrackReaderToolCfg(flags)))
        kwargs.setdefault("GNNTrackFinderTool", None)
    else:
        raise RuntimeError("GNNTrackFinder or GNNTrackReader must be enabled!")
 
    kwargs.setdefault("SeedFitterTool", acc.popToolsAndMerge(SeedFitterToolCfg(flags)))
 
    from TrkConfig.CommonTrackFitterConfig import ITkTrackFitterCfg
    kwargs.setdefault("TrackFitter", acc.popToolsAndMerge(ITkTrackFitterCfg(flags)))
 
    from InDetConfig.SiDetElementsRoadToolConfig import ITkSiDetElementsRoadMaker_xkCfg
    kwargs.setdefault("RoadTool", acc.popToolsAndMerge(ITkSiDetElementsRoadMaker_xkCfg(flags)))
 
    # configurations for Kalman filter.
    # similar to https://gitlab.cern.ch/atlas/athena/-/blob/main/InnerDetector/InDetConfig/python/SiTrackMakerConfig.py#L188
    kwargs.setdefault("nClustersMin", flags.Tracking.ActiveConfig.minClusters[0])
    kwargs.setdefault("nWeightedClustersMin", flags.Tracking.ActiveConfig.nWeightedClustersMin[0])
    kwargs.setdefault("nHolesMax", flags.Tracking.ActiveConfig.nHolesMax[0])
    kwargs.setdefault("nHolesGapMax", flags.Tracking.ActiveConfig.nHolesGapMax[0])
 
    kwargs.setdefault("pTmin", flags.Tracking.ActiveConfig.minPT[0])
    kwargs.setdefault("pTminBrem", flags.Tracking.ActiveConfig.minPTBrem[0])
    kwargs.setdefault("Xi2max", flags.Tracking.ActiveConfig.Xi2max[0])
    kwargs.setdefault("Xi2maxNoAdd", flags.Tracking.ActiveConfig.Xi2maxNoAdd[0])
    kwargs.setdefault("Xi2maxMultiTracks", flags.Tracking.ActiveConfig.Xi2max[0])
    kwargs.setdefault("doMultiTracksProd", False)
 
    acc.addEventAlgo(CompFactory.InDet.GNNSeedingTrackMaker(name, **kwargs))
    return acc