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 filter_segment_container = (
18 args.filterSegmentsWithoutMlConnections and run_edge_classifier
19 )
20 filtered_segment_key = "MuonSegmentsFromR4MlConnected"
21
22 if args.skip_onnx and (args.enableBucketFilter or args.enableEdgeClassifier):
23 print(
"INFO: --skip-onnx requested. Disabling bucket filter and edge classifier inference stages.")
24 if args.skip_onnx and args.useMlSeeder:
25 print(
"INFO: --skip-onnx requested. Switching to the standard seeder for the non-ONNX baseline.")
26 if args.filterSegmentsWithoutMlConnections and not run_edge_classifier:
27 print(
"WARNING: --filterSegmentsWithoutMlConnections requires the edge "
28 "classifier and will be ignored.")
29
30 if args.athenaDebug:
31 flags.Exec.DebugMessageComponents = [
32 "GraphInferenceAlg",
33 "GraphInferenceAlg.GraphBucketFilterTool",
34 "GraphInferenceAlg.GraphBucketFilterTool.OnnxRuntimeSessionToolCPU",
35 "GraphInferenceAlg.GraphBucketFilterTool.OnnxRuntimeSessionToolCUDA",
36 "SegmentEdgeInferenceAlg",
37 "SegmentEdgeInferenceAlg.SegmentEdgeClassifierTool",
38 "SegmentEdgeInferenceAlg.SegmentEdgeClassifierTool.OnnxRuntimeSessionToolCPU",
39 "SegmentEdgeInferenceAlg.SegmentEdgeClassifierTool.OnnxRuntimeSessionToolCUDA",
40 "MSTrackFinderAlg",
41 "MSTrackFinderAlg.MlMsTrackSeeder",
42 ]
43
44 from AthOnnxComps.OnnxRuntimeFlags import OnnxRuntimeType
45 if run_bucket_filter or run_edge_classifier:
46 flags.AthOnnx.ExecutionProvider = (
47 OnnxRuntimeType.CPU if args.use_cpu else OnnxRuntimeType.CUDA
48 )
49 else:
50 flags.AthOnnx.ExecutionProvider = OnnxRuntimeType.CPU
51
52 flags, cfg = setupGeoR4TestCfg(args, flags)
53
54 if not args.skipTrackTester:
55 cfg.merge(setupHistSvcCfg(flags, outFile=args.outRootFile,
56 outStream="MuonTrackTester"))
57
58 output_level = 1 if args.athenaDebug else 3
59
60 if run_bucket_filter:
61 from MuonInference.InferenceConfig import GraphBucketFilterToolCfg, GraphInferenceAlgCfg
62 bucketTool = cfg.popToolsAndMerge(
63 GraphBucketFilterToolCfg(
64 flags,
65 ModelPath=args.bucket_model_path,
66 ScoreThreshold=args.score_threshold,
67 OutputName=args.output_name,
68 SingleOutputMode=args.single_output_mode,
69 OutputLevel=output_level,
70 )
71 )
72 cfg.merge(GraphInferenceAlgCfg(flags, InferenceTools=[bucketTool]))
73
74 from MuonConfig.ReconstructionConfigR4 import MuonReconstructionConfig
76 if run_bucket_filter:
77 cfg.getEventAlgo("MuonEtaHoughTransformAlg").SpacePointContainer = "FilteredMlBuckets"
78
79 if run_edge_classifier:
80 from MuonInference.InferenceConfig import SegmentEdgeInferenceAlgCfg
81 edge_classifier_kwargs = {
82 "ModelPath": args.edgeModel,
83 "ReadSpacePoints": (
84 "FilteredMlBuckets" if run_bucket_filter else "MuonSpacePoints"
85 ),
86
87
88 "MaxSegmentsPerBucket": args.maxSegmentsPerBucket,
89 "MaxEdgesPerNodeBeforeInference": args.maxEdgesBeforeInference,
90 "MaxEdgesPerTargetChamberBeforeInference": (
91 args.maxEdgesPerTargetChamber
92 ),
93 "DropSameChamberEdgesBeforeInference": (
94 not args.keepSameChamberEdgesBeforeInference
95 ),
96 "DropIsolatedNodesBeforeInference": (
97 not args.keepIsolatedNodesBeforeInference
98 ),
99 }
100 if args.maxDeltaThetaDeg is not None:
101 edge_classifier_kwargs["MaxDeltaThetaDeg"] = args.maxDeltaThetaDeg
102 edge_inference_kwargs = {
103 "EdgeClassifierTool": edge_classifier_kwargs,
104 "PairGateDecoration": "MuonSegmentsFromR4.mlTrackComponent",
105 "PairGateThreshold": args.edgeThreshold,
106 "MaxEdgesPerNode": args.maxEdgesPerSegment,
107 "UseDegreeCappedComponents": args.useDegreeCappedMlComponents,
108 "RequireMutualTopKEdges": not args.allowOneSidedMlEdges,
109 "RecoverOrphanNodes": not args.disableOrphanRecovery,
110 "SeedAnchorsPerComponent": args.seedAnchorsPerComponent,
111 "MinSegmentsPerComponent": args.minSegmentsPerComponent,
112 "KeepBestSegmentPerChamber": not args.keepAllSegmentsPerChamber,
113 "OutputLevel": output_level,
114 }
115 if filter_segment_container:
116 edge_inference_kwargs["FilteredSegmentKey"] = filtered_segment_key
117 cfg.merge(SegmentEdgeInferenceAlgCfg(flags, **edge_inference_kwargs))
118
119 if run_ml_seeder and not run_edge_classifier:
120 print(
"WARNING: ML seeder enabled while edge classifier is disabled."
121 " The decoration 'mlTrackComponent' may be missing.")
122
123 ms_track_finder = cfg.getEventAlgo("MSTrackFinderAlg")
124 ms_track_finder.OutputLevel = output_level
125 ms_track_finder.UseMlSeeder = run_ml_seeder
126 ms_track_finder.MlCandidateDecoration = "mlTrackComponent"
127 if filter_segment_container:
128
129 ms_track_finder.SegmentContainer = filtered_segment_key
130
131 ms_track_finder.MlCandidateDecorationKey = (
132 "MuonSegmentsFromR4.mlTrackComponent"
133 )
134 ms_track_finder.MlFallbackToBaselineIfUndecorated = True
135 ms_track_finder.MlFallbackToBaselineIfNoCandidates = False
136
137 if not args.skipTrackTester:
138 from MuonTrackFindingTest.MsTrackFindingTester import MsTrackTesterCfg
139 cfg.merge(MsTrackTesterCfg(flags, scheduleLegacy=False, outFile=args.outRootFile))
140
141 if args.enableRecoChainTester:
142 from MuonTrackFindingAlgs.TrackFindingConfig import MuonActsToTrkConvCfg
143 cfg.merge(MuonActsToTrkConvCfg(flags,
144 ACTSTracksLocation="MsTracks",
145 TracksLocation="MsTracksConv"))
146
147 from xAODTrackingCnv.xAODTrackingCnvConfig import MuonStandaloneTrackParticleCnvAlgCfg
148 cfg.merge(MuonStandaloneTrackParticleCnvAlgCfg(flags,
149 name="MuonXAODParticleConvR4",
150 TrackContainerName="MsTracksConv",
151 xAODTrackParticlesFromTracksContainerName="MuonSpectrometerTrackParticlesR4"))
152
153 if flags.Input.isMC:
154 from MuonTruthAlgsR4.MuonTruthAlgsConfig import RecoSegmentTruthAssocCfg, TrackToTruthPartAssocCfg
155 cfg.merge(RecoSegmentTruthAssocCfg(flags,
156 name="MuonSegmentsFromR4TruthMatching",
157 SegmentKey="MuonSegmentsFromR4"))
158 cfg.merge(TrackToTruthPartAssocCfg(flags,
159 name="TrackToTruthMuonSpectrometerTrackParticlesR4",
160 TrackCollection="MuonSpectrometerTrackParticlesR4"))
161
162 from MuonPatternRecognitionTest.PatternTestConfig import MuonRecoChainTesterCfg
163 cfg.merge(MuonRecoChainTesterCfg(flags,
164 LegacySegmentKey="MuonSegmentsFromR4",
165 SegmentFromR4HoughKey="",
166 R4SegmentKey="MuonSegmentsFromR4",
167 LegacyTrackKey="MuonSpectrometerTrackParticlesR4",
168 TrackKeyHoughR4="",
169 TrackKeyR4="MuonSpectrometerTrackParticlesR4"))
170
171 cfg.printConfig(withDetails=True, summariseProps=True)
172 executeTest(cfg)
173
void print(char *figname, TCanvas *c1)
1.17.0