4 from MuonGeoModelTestR4.testGeoModel import setupGeoR4TestCfg
5 from MuonConfig.MuonConfigUtils import executeTest, setupHistSvcCfg
6 from AthenaConfiguration.AllConfigFlags import initConfigFlags
7 flags = initConfigFlags()
8 flags.PerfMon.doFullMonMT = False
9 includeG4TrackTruth = bool(getattr(args, "includeG4TrackTruth", False))
10 minG4TrackTruthSegments = int(getattr(args, "minG4TrackTruthSegments", 2))
11 if minG4TrackTruthSegments < 1:
12 raise ValueError("--minG4TrackTruthSegments must be >= 1")
13 if includeG4TrackTruth:
14 flags.Muon.includePileUpTruth = True
15
16 from AthOnnxComps.OnnxRuntimeFlags import OnnxRuntimeType
17
18 if getattr(args, "use_cpu", False):
19 flags.AthOnnx.ExecutionProvider = OnnxRuntimeType.CPU
20 else:
21 flags.AthOnnx.ExecutionProvider = OnnxRuntimeType.CUDA
22
23 flags, cfg = setupGeoR4TestCfg(args)
24
25 cfg.merge(setupHistSvcCfg(flags, outFile=args.outRootFile, outStream="MuonSegmentDump"))
26
27 from MuonConfig.MuonDataPrepConfig import xAODUncalibMeasPrepCfg
28 cfg.merge(xAODUncalibMeasPrepCfg(flags))
29
30 from MuonSpacePointFormation.SpacePointFormationConfig import MuonSpacePointFormationCfg
31 cfg.merge(MuonSpacePointFormationCfg(flags))
32
33 from MuonPatternRecognitionAlgs.MuonPatternRecognitionConfig import MuonPatternRecognitionCfg
34
35 do_ml_bucket_filter = bool(getattr(args, "doMLBucketFilter", False) or
36 getattr(args, "bucketModel", None) is not None or
37 getattr(args, "bucketThreshold", None) is not None)
38 if do_ml_bucket_filter:
39 from MuonInference.InferenceConfig import GraphBucketFilterToolCfg, GraphInferenceAlgCfg
40 bucket_tool_kwargs: dict[str, object] = {"WriteSpacePointKey": "FilteredMlBuckets"}
41 if getattr(args, "bucketModel", None) is not None:
42 bucket_tool_kwargs["ModelPath"] = args.bucketModel
43 if getattr(args, "bucketThreshold", None) is not None:
44 bucket_tool_kwargs["ScoreThreshold"] = args.bucketThreshold
45 bucket_tool_kwargs["OutputName"] = args.output_name
46 bucket_tool_kwargs["SingleOutputIsLogit"] = args.is_logit if hasattr(args, "is_logit") else False
47 bucket_tool = cfg.popToolsAndMerge(
48 GraphBucketFilterToolCfg(
49 flags,
50 **bucket_tool_kwargs,
51 )
52 )
53 cfg.merge(
54 GraphInferenceAlgCfg(
55 flags,
56 InferenceTools=[bucket_tool],
57 )
58 )
59
60
61 cfg.merge(MuonPatternRecognitionCfg(flags))
62 cfg.getEventAlgo("MuonEtaHoughTransformAlg").SpacePointContainer = "FilteredMlBuckets"
63 else:
64 cfg.merge(MuonPatternRecognitionCfg(flags))
65
66
67 if flags.Input.isMC:
68 from MuonTruthAlgsR4.MuonTruthAlgsConfig import MuonTruthAlgsCfg
69 cfg.merge(MuonTruthAlgsCfg(flags))
70
71 from MuonBucketDump.MuonBucketDumpConfig import MuonSegmentDumpCfg
72 dumper_kwargs = {
73 "IncludeG4TrackTruth": includeG4TrackTruth,
74 "MinG4TrackTruthSegments": minG4TrackTruthSegments,
75 }
76
77 if do_ml_bucket_filter:
78 cfg.merge(MuonSegmentDumpCfg(flags, SpacePointKeys=["FilteredMlBuckets"],
79 **dumper_kwargs))
80 else:
81 cfg.merge(MuonSegmentDumpCfg(flags, **dumper_kwargs))
82
83 executeTest(cfg)
84
85
1.16.1