6 from MuonGeoModelTestR4.testGeoModel import setupGeoR4TestCfg
7 from MuonConfig.MuonConfigUtils import executeTest, setupHistSvcCfg
8 from AthenaConfiguration.AllConfigFlags import initConfigFlags
9 flags = initConfigFlags()
10 flags.PerfMon.doFullMonMT = not args.noPerfMon
11 flags.PerfMon.OutputJSON = "perfmonmt_MuonR4Reco.json"
12 flags.Trigger.Muon.useNewRegionSelector = False
13
14 run_bucket_filter = args.enableBucketFilter and not args.skip_onnx
15 run_edge_classifier = args.enableEdgeClassifier and not args.skip_onnx
16 run_ml_seeder = args.useMlSeeder and not args.skip_onnx
17
18 if args.skip_onnx and (args.enableBucketFilter or args.enableEdgeClassifier):
19 print(
"INFO: --skip-onnx requested. Disabling bucket filter and edge classifier inference stages.")
20 if args.skip_onnx and args.useMlSeeder:
21 print(
"INFO: --skip-onnx requested. Switching to legacy seeder for a non-ONNX baseline.")
22
23 if args.athenaDebug:
24 flags.Exec.DebugMessageComponents = [
25 "GraphInferenceAlg",
26 "GraphInferenceAlg.GraphBucketFilterTool",
27 "GraphInferenceAlg.GraphBucketFilterTool.OnnxRuntimeSessionToolCPU",
28 "GraphInferenceAlg.GraphBucketFilterTool.OnnxRuntimeSessionToolCUDA",
29 "SegmentEdgeInferenceAlg",
30 "SegmentEdgeInferenceAlg.SegmentEdgeClassifierTool",
31 "SegmentEdgeInferenceAlg.SegmentEdgeClassifierTool.OnnxRuntimeSessionToolCPU",
32 "SegmentEdgeInferenceAlg.SegmentEdgeClassifierTool.OnnxRuntimeSessionToolCUDA",
33 "SegmentEdgeInferenceAlg.SegmentTrackCandidateBuilderTool",
34 "MSTrackFinderAlg",
35 "MSTrackFinderAlg.MlMsTrackSeeder",
36 ]
37
38 from AthOnnxComps.OnnxRuntimeFlags import OnnxRuntimeType
39 if run_bucket_filter or run_edge_classifier:
40 flags.AthOnnx.ExecutionProvider = (
41 OnnxRuntimeType.CPU if args.use_cpu else OnnxRuntimeType.CUDA
42 )
43 else:
44 flags.AthOnnx.ExecutionProvider = OnnxRuntimeType.CPU
45
46 flags, cfg = setupGeoR4TestCfg(args, flags)
47
48 if not args.skipTrackTester:
49 cfg.merge(setupHistSvcCfg(flags, outFile=args.outRootFile,
50 outStream="MuonTrackTester"))
51
52 output_level = 1 if args.athenaDebug else 3
53
54 if run_bucket_filter:
55 from MuonInference.InferenceConfig import GraphBucketFilterToolCfg, GraphInferenceAlgCfg
56 bucketTool = cfg.popToolsAndMerge(
57 GraphBucketFilterToolCfg(
58 flags,
59 ModelPath=args.bucket_model_path,
60 ScoreThreshold=args.score_threshold,
61 OutputName=args.output_name,
62 SingleOutputMode=args.single_output_mode,
63 OutputLevel=output_level,
64 )
65 )
66 cfg.merge(GraphInferenceAlgCfg(flags, InferenceTools=[bucketTool]))
67
68 from MuonConfig.ReconstructionConfigR4 import MuonReconstructionConfig
70 if run_bucket_filter:
71 cfg.getEventAlgo("MuonEtaHoughTransformAlg").SpacePointContainer = "FilteredMlBuckets"
72
73 if run_edge_classifier:
74 from MuonInference.InferenceConfig import SegmentEdgeInferenceAlgCfg
75 edge_classifier_kwargs = {
76 "ModelPath": args.edgeModel,
77 "ReadSpacePoints": (
78 "FilteredMlBuckets" if run_bucket_filter else "MuonSpacePoints"
79 ),
80 }
81 cfg.merge(SegmentEdgeInferenceAlgCfg(
82 flags,
83 EdgeClassifierTool=edge_classifier_kwargs,
84 EdgeThreshold=args.edgeThreshold,
85 OverlapThreshold=args.overlapThreshold,
86 UseRecoveryComponents=args.useRecoveryComponents,
87 OutputLevel=output_level,
88 ))
89
90 if run_ml_seeder and not run_edge_classifier:
91 print(
"WARNING: ML seeder enabled while edge classifier is disabled."
92 " The decoration 'trackCandidateIds' may be missing.")
93
94 ms_track_finder = cfg.getEventAlgo("MSTrackFinderAlg")
95 ms_track_finder.UseMlSeeder = run_ml_seeder
96 ms_track_finder.MlCandidateDecoration = "trackCandidateIds"
97 ms_track_finder.MlFallbackToBaselineIfUndecorated = True
98 ms_track_finder.MlFallbackToBaselineIfNoCandidates = True
99
100 if not args.skipTrackTester:
101 from MuonTrackFindingTest.MsTrackFindingTester import MsTrackTesterCfg
102 cfg.merge(MsTrackTesterCfg(flags, scheduleLegacy=False, outFile=args.outRootFile))
103
104 if args.enableRecoChainTester:
105 from MuonTrackFindingAlgs.TrackFindingConfig import MuonActsToTrkConvCfg
106 cfg.merge(MuonActsToTrkConvCfg(flags,
107 ACTSTracksLocation="MsTracks",
108 TracksLocation="MsTracksConv"))
109
110 from xAODTrackingCnv.xAODTrackingCnvConfig import MuonStandaloneTrackParticleCnvAlgCfg
111 cfg.merge(MuonStandaloneTrackParticleCnvAlgCfg(flags,
112 name="MuonXAODParticleConvR4",
113 TrackContainerName="MsTracksConv",
114 xAODTrackParticlesFromTracksContainerName="MuonSpectrometerTrackParticlesR4"))
115
116 if flags.Input.isMC:
117 from MuonTruthAlgsR4.MuonTruthAlgsConfig import RecoSegmentTruthAssocCfg, TrackToTruthPartAssocCfg
118 cfg.merge(RecoSegmentTruthAssocCfg(flags,
119 name="MuonSegmentsFromR4TruthMatching",
120 SegmentKey="MuonSegmentsFromR4"))
121 cfg.merge(TrackToTruthPartAssocCfg(flags,
122 name="TrackToTruthMuonSpectrometerTrackParticlesR4",
123 TrackCollection="MuonSpectrometerTrackParticlesR4"))
124
125 from MuonPatternRecognitionTest.PatternTestConfig import MuonRecoChainTesterCfg
126 cfg.merge(MuonRecoChainTesterCfg(flags,
127 LegacySegmentKey="MuonSegmentsFromR4",
128 SegmentFromR4HoughKey="",
129 R4SegmentKey="MuonSegmentsFromR4",
130 LegacyTrackKey="MuonSpectrometerTrackParticlesR4",
131 TrackKeyHoughR4="",
132 TrackKeyR4="MuonSpectrometerTrackParticlesR4"))
133
134 cfg.printConfig(withDetails=True, summariseProps=True)
135 executeTest(cfg)
136
void print(char *figname, TCanvas *c1)