ATLAS Offline Software
Loading...
Searching...
No Matches
TrigTauRecToolsConfig.py
Go to the documentation of this file.
1# Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
2
3from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
4from AthenaConfiguration.ComponentFactory import CompFactory
5from AthenaCommon.SystemOfUnits import mm
6
7def trigTauVertexFinderCfg(flags, name=''):
8 '''Algorithm that overwrites numTrack() and charge() of tauJets in container'''
9 acc = ComponentAccumulator()
10
11 acc.setPrivateTools(CompFactory.TauVertexFinder(
12 name = name,
13 UseTJVA = False,
14 AssociatedTracks = 'GhostTrack',
15 InDetTrackSelectionToolForTJVA = '',
16 Key_vertexInputContainer = '',
17 TVATool = '',
18 ))
19
20 return acc
21
22def trigTauTrackFinderCfg(flags, name='', TrackParticlesContainer=''):
23 '''Tau track association'''
24 acc = ComponentAccumulator()
25
26 from TrkConfig.TrkVertexFitterUtilsConfig import AtlasTrackToVertexIPEstimatorCfg
27 AtlasTrackToVertexIPEstimator = acc.popToolsAndMerge(AtlasTrackToVertexIPEstimatorCfg(flags))
28
29 from TrackToVertex.TrackToVertexConfig import TrackToVertexCfg
30 TrackToVertexTool = acc.popToolsAndMerge(TrackToVertexCfg(flags))
31
32 from InDetConfig.InDetTrackSelectorToolConfig import TrigTauInDetTrackSelectorToolCfg
33 TrigTauInDetTrackSelectorTool = acc.popToolsAndMerge(TrigTauInDetTrackSelectorToolCfg(flags))
34
35 acc.setPrivateTools(CompFactory.TauTrackFinder(
36 name = name,
37 MaxJetDrTau = 0.2,
38 MaxJetDrWide = 0.4,
39 TrackSelectorToolTau = TrigTauInDetTrackSelectorTool,
40 TrackToVertexTool = TrackToVertexTool,
41 Key_trackPartInputContainer = TrackParticlesContainer,
42 maxDeltaZ0wrtLeadTrk = 0.75*mm,
43 TrackToVertexIPEstimator = AtlasTrackToVertexIPEstimator,
44 removeDuplicateCoreTracks = True,
45 ))
46
47 return acc
48
49def tauVertexVariablesCfg(flags, name=''):
50 '''Vertex variables calculation'''
51 acc = ComponentAccumulator()
52
53 from TrkConfig.TrkVertexFittersConfig import TauAdaptiveVertexFitterCfg
54 TauAdaptiveVertexFitter = acc.popToolsAndMerge(TauAdaptiveVertexFitterCfg(flags))
55
56 from TrkConfig.TrkVertexSeedFinderToolsConfig import CrossDistancesSeedFinderCfg
57 CrossDistancesSeedFinder = acc.popToolsAndMerge(CrossDistancesSeedFinderCfg(flags))
58
59 acc.setPrivateTools(CompFactory.TauVertexVariables(
60 name = name,
61 VertexFitter = TauAdaptiveVertexFitter,
62 SeedFinder = CrossDistancesSeedFinder,
63 ))
64
65 return acc
66
67def trigTauJetLVNNEvaluatorCfg(flags, tau_id='', use_taujet_rnnscore=True):
68 '''TauJet identification inference based on LVNN models, for RNNs and DeepSets'''
69 acc = ComponentAccumulator()
70
71 try: id_flags = getattr(flags.Trigger.Offline.Tau, tau_id)
72 except NameError: raise ValueError(f'Invalid LVNN (RNN, DeepSet) TauID configuration: {tau_id}')
73
74 # For legacy mediumRNN/tightRNN_tracktwoMVA/tracktwoLLP/trackLRT chains
75 output_variable = 'RNNJetScore' if use_taujet_rnnscore else f'{tau_id}_Score'
76
77 # This is to prevent clashes between the same ID running over the legacy chains, and the new decorator-based chains
78 sfx = '_RNNJetScore' if use_taujet_rnnscore else ''
79
80 acc.setPrivateTools(CompFactory.TauJetRNNEvaluator(
81 name = f'TrigTau_TauJetLVNNEvaluator_{tau_id}{sfx}',
82 useTRT = flags.Detector.EnableTRT,
83
84 # Network config:
85 NetworkFile0P = id_flags.NetworkConfig[0],
86 NetworkFile1P = id_flags.NetworkConfig[1],
87 NetworkFile3P = id_flags.NetworkConfig[2],
88 InputLayerScalar = 'scalar',
89 InputLayerTracks = 'tracks',
90 InputLayerClusters = 'clusters',
91 OutputLayer = 'rnnid_output',
92 OutputNode = 'sig_prob',
93
94 # Inputs:
95 MaxTracks = id_flags.MaxTracks,
96 MaxClusters = id_flags.MaxClusters,
97 MaxClusterDR = 1.0,
98 VertexCorrection = False,
99 TrackClassification = False,
100
101 # Decorated TauJet variable names:
102 OutputVarname = output_variable,
103 ))
104
105 return acc
106
107def trigTauJetONNXEvaluatorCfg(
108 flags,
109 tau_id: str,
110 tau_container: str = '',
111 hits_decoration_container: str = '',
112 ptau_sfx: str | None = None,
113 pjet_sfx: str | None = None,
114):
115 '''TauJet identification inference based on ONNX models, for GNNs, transformers, etc...'''
116 acc = ComponentAccumulator()
117
118 try: id_flags = getattr(flags.Trigger.Offline.Tau, tau_id)
119 except NameError: raise ValueError(f'Invalid ONNX TauID configuration: {tau_id}')
120
121 if isinstance(id_flags.ONNXConfig, str):
122 network_config = {'NetworkFileInclusive': id_flags.ONNXConfig}
123 elif isinstance(id_flags.ONNXConfig, (list, tuple)) and len(id_flags.ONNXConfig) == 3:
124 network_config = {
125 'NetworkFile0P': id_flags.ONNXConfig[0],
126 'NetworkFile1P': id_flags.ONNXConfig[1],
127 'NetworkFile3P': id_flags.ONNXConfig[2],
128 }
129 else:
130 raise ValueError(f'Invalid {tau_id} ONNX network config file')
131
132 name = f'TrigTau_TauJetONNXEvaluator_{tau_id}'
133 if tau_container: name += f'_{tau_container}'
134 if hits_decoration_container: name += f'_{hits_decoration_container}'
135
136 ptau_decor = id_flags.OutputPTau if hasattr(id_flags, 'OutputPTau') else f'{tau_id}_ProbTau'
137 if ptau_sfx: ptau_decor += f'_{ptau_sfx}'
138 pjet_decor = id_flags.OutputPJet if hasattr(id_flags, 'OutputPJet') else f'{tau_id}_ProbJet'
139 if pjet_sfx: pjet_decor += f'_{pjet_sfx}'
140
141 acc.setPrivateTools(CompFactory.TauGNNEvaluator(
142 name = name,
143 useTRT = flags.Detector.EnableTRT,
144
145 # Network config:
146 **network_config,
147 InputLayerScalar = 'tau_vars',
148 InputLayerTracks = 'track_vars',
149 InputLayerClusters = 'cluster_vars',
150 InputLayerHits = 'hit_vars',
151 NodeNameTau = id_flags.NodeNameTau if hasattr(id_flags, 'NodeNameTau') else 'pTau',
152 NodeNameJet = id_flags.NodeNameJet if hasattr(id_flags, 'NodeNameJet') else 'pJet',
153
154 # Inputs:
155 MaxTracks = id_flags.MaxTracks if hasattr(id_flags, 'MaxTracks') else 0,
156 MaxClusters = id_flags.MaxClusters if hasattr(id_flags, 'MaxClusters') else 0,
157 MaxClusterDR = 1.0,
158 MaxHits = id_flags.MaxHits if hasattr(id_flags, 'MaxHits') else 0,
159 TauContainerName = tau_container,
160 HitsHandleKey = hits_decoration_container,
161 VertexCorrection = False,
162 TrackClassification = False,
163
164 # Decorated TauJet variable names:
165 OutputVarname = f'{tau_id}_Score',
166 OutputDiscriminant = id_flags.OutputDiscriminant,
167 OutputPTau = ptau_decor,
168 OutputPJet = pjet_decor,
169
170 # Tau prongness selection
171 MinProngTrackPt = id_flags.MinProngTrackPt if hasattr(id_flags, 'MinProngTrackPt') else 0,
172 ))
173
174 return acc
175
176def trigTauWPDecoratorRNNCfg(flags, tau_id: str, precision_seq_name: str):
177 '''
178 TauJet signal transformed score and ID WPs decorator tool,
179 for the legacy mediumRNN/tightRNN_tracktwoMVA/tracktwoLLP/trackLRT chains ONLY!
180 '''
181
182 acc = ComponentAccumulator()
183
184 try: id_flags = getattr(flags.Trigger.Offline.Tau, tau_id)
185 except NameError: raise ValueError(f'Invalid TauID configuration: {tau_id}')
186
187 # In this version we store the WPs as flags accessable through
188 # tau->isTau(xAOD::TauJetParameters::IsTauFlags::JetRNNSig...)
189 import PyUtils.RootUtils as ru
190 ROOT = ru.import_root()
191 import cppyy
192 cppyy.load_library('libxAODTau_cDict')
193 WPEnumVals = [
194 ROOT.xAOD.TauJetParameters.IsTauFlag.JetRNNSigVeryLoose,
195 ROOT.xAOD.TauJetParameters.IsTauFlag.JetRNNSigLoose,
196 ROOT.xAOD.TauJetParameters.IsTauFlag.JetRNNSigMedium,
197 ROOT.xAOD.TauJetParameters.IsTauFlag.JetRNNSigTight,
198 ]
199
200 acc.setPrivateTools(CompFactory.TauWPDecorator(
201 name=f'TrigTau_TauWPDecoratorRNN_{precision_seq_name}_{tau_id}',
202 flatteningFile0Prong=id_flags.ScoreFlatteningConfig[0],
203 flatteningFile1Prong=id_flags.ScoreFlatteningConfig[1],
204 flatteningFile3Prong=id_flags.ScoreFlatteningConfig[2],
205 CutEnumVals=WPEnumVals,
206 SigEff0P=id_flags.TargetEff[0],
207 SigEff1P=id_flags.TargetEff[1],
208 SigEff3P=id_flags.TargetEff[2],
209 ScoreName='RNNJetScore',
210 NewScoreName='RNNJetScoreSigTrans',
211 DefineWPs=True,
212 ))
213
214 return acc
215
216def trigTauWPDecoratorCfg(flags, tau_id: str, precision_seq_name: str, tauContainerName: str):
217 '''TauJet signal transformed score and ID WPs decorator tool'''
218 acc = ComponentAccumulator()
219
220 try: id_flags = getattr(flags.Trigger.Offline.Tau, tau_id)
221 except NameError: raise ValueError(f'Invalid TauID configuration: {tau_id}')
222
223 if isinstance(id_flags.ScoreFlatteningConfig, str):
224 cfg = {
225 'flatteningFile0Prong': id_flags.ScoreFlatteningConfig,
226 'flatteningFile1Prong': id_flags.ScoreFlatteningConfig,
227 'flatteningFile3Prong': id_flags.ScoreFlatteningConfig,
228 }
229 elif isinstance(id_flags.ScoreFlatteningConfig, (list, tuple)) and len(id_flags.ScoreFlatteningConfig) == 3:
230 cfg = {
231 'flatteningFile0Prong': id_flags.ScoreFlatteningConfig[0],
232 'flatteningFile1Prong': id_flags.ScoreFlatteningConfig[1],
233 'flatteningFile3Prong': id_flags.ScoreFlatteningConfig[2],
234 }
235 else:
236 raise ValueError(f'Invalid {tau_id} WP decorator flattening config')
237
238
239 acc.setPrivateTools(CompFactory.TauWPDecorator(
240 name=f'TrigTau_TauWPDecorator_{precision_seq_name}_{tau_id}',
241 **cfg,
242 TauContainerName=tauContainerName,
243 DecorWPNames=[f'{tau_id}_{wp}' for wp in id_flags.TargetWPs],
244 DecorWPCutEffs0P=[eff[0] for eff in id_flags.TargetWPs.values()],
245 DecorWPCutEffs1P=[eff[1] for eff in id_flags.TargetWPs.values()],
246 DecorWPCutEffs3P=[eff[2] for eff in id_flags.TargetWPs.values()],
247 ScoreName=f'{tau_id}_Score',
248 NewScoreName=f'{tau_id}_ScoreSigTrans',
249 DefineWPs=True,
250 ))
251
252 return acc
trigTauVertexFinderCfg(flags, name='')