ATLAS Offline Software
Loading...
Searching...
No Matches
python.OutputAnalysisConfig.OutputAnalysisConfig Class Reference
Inheritance diagram for python.OutputAnalysisConfig.OutputAnalysisConfig:
Collaboration diagram for python.OutputAnalysisConfig.OutputAnalysisConfig:

Public Member Functions

 __init__ (self)
 instanceName (self)
 createOutputAlgs (self, config, name, vars, nonContainers)
 makeAlgs (self, config)
 createSelectionFlagBranches (self, config, containers)
 makeSelectionSummaryAlg (self, config, containerName, selectionName)

Static Public Member Functions

 branchSortOrder (rule)

Public Attributes

 containersFullMET
 containersOnlyForMC
 containersOnlyForDSIDs
 varsOnlyForDSIDs
 storeSelectionFlags
 alwaysAddNosys
 metTermName
 truthMetTermName
str outputFormat = 'RNTuple':

Detailed Description

the ConfigBlock for the output ntuple configuration

Definition at line 10 of file OutputAnalysisConfig.py.

Constructor & Destructor Documentation

◆ __init__()

python.OutputAnalysisConfig.OutputAnalysisConfig.__init__ ( self)

Definition at line 13 of file OutputAnalysisConfig.py.

13 def __init__ (self) :
14 super().__init__ ()
15 self.addOption ('postfix', '', type=str,
16 info="a postfix to apply to decorations and algorithm names. "
17 "Typically not needed here.")
18 self.addOption ('vars', [], type=list,
19 info="a list of mappings (list of strings) between containers and "
20 "decorations to output branches.")
21 self.addOption ('varsOnlyForMC', [], type=list,
22 info="same as `vars`, but for MC-only variables so as to avoid a "
23 "crash when running on data.")
24 self.addOption ('metVars', [], type=list,
25 info="a list of mappings (list of strings) between containers "
26 "and decorations to output branches. Specficially for MET "
27 "variables, where only the final MET term is retained.")
28 self.addOption ('truthMetVars', [], type=list,
29 info="a list of mappings (list of strings) between containers "
30 "and decorations to output branches for truth MET.")
31 self.addOption ('containers', {}, type=dict,
32 info="a dictionary mapping prefixes (key) to container names "
33 "(values) to be used when saving to the output tree. Branches "
34 "are then of the form `prefix_decoration`.")
35 self.addOption ('containersFullMET', {}, type=dict,
36 info="same as `containers`, but for MET containers that should be "
37 "saved with all terms (as opposed to just the final term). This "
38 "is useful for special studies. A container can appear both here and "
39 "in containers (with different prefixes).")
40 self.addOption ('containersOnlyForMC', {}, type=dict,
41 info="same as `containers`, but for MC-only containers so as to avoid "
42 "a crash when running on data.")
43 self.addOption ('containersOnlyForDSIDs', {}, type=dict,
44 info="specify which DSIDs are allowed to produce a given container. "
45 "This works like `onlyForDSIDs`: pass a list of DSIDs or regexps.")
46 self.addOption ('varsOnlyForDSIDs', {}, type=dict,
47 info="specify which DSIDs are allowed to produce a given variable. "
48 "Pass a dictionary with the variable definition as key and a list of DSIDs/regexps as value.")
49 self.addOption ('nonContainers', [], type=list,
50 info="a list of container names that are not actual containers but should be treated as non-containers.")
51 self.addOption ('treeName', 'analysis', type=str,
52 info="name of the output TTree (or RNTuple) to save.")
53 self.addOption ('streamName', 'ANALYSIS', type=str,
54 info="name of the output stream to save the tree in.")
55 self.addOption ('metTermName', 'Final', type=str,
56 info="the name of the MET term to save, turning the MET "
57 "container into a single object.")
58 self.addOption ('truthMetTermName', 'NonInt', type=str,
59 info="the name of the truth MET term to save, turning the MET "
60 "container into a single object.")
61 self.addOption ('storeSelectionFlags', True, type=bool,
62 info="whether to store one branch for each object selection.")
63 self.addOption ('selectionFlagPrefix', 'select', type=str,
64 info="the prefix used when naming selection branches.")
65 self.addOption ('commands', [], type=list,
66 info="a list of strings containing commands (regexp strings "
67 "prefaced by the keywords `enable` or `disable`) to turn on/off the "
68 "writing of branches to the output ntuple. If left empty, do not modify "
69 "the scheduled output branches.")
70 self.addOption ('commandsOnlyForDSIDs', {}, type=dict,
71 info="a dictionary with individual DSIDs as keys, and a list of strings "
72 "like for the `commands` option as items. These `commands` will only be run "
73 "for the corresponding DSID.")
74 self.addOption ('alwaysAddNosys', False, type=bool,
75 info="If set to `True`, all branches will be given a systematics suffix, "
76 "even if they have no systematics (beyond the nominal).")
77 self.addOption ('skipRedundantSelectionFlags', True, type=bool,
78 info="remove the redundant 'outputSelect' branches created by the Thinning step. "
79 "These could however be used to simplify downstream workflows, as in Easyjet. "
80 "The default is True.")
81 self.addOption ('outputFormat', 'TTree', type=str,
82 info="The output format, `TTree` or `RNTuple`.")
83 self.addOption ('defaultBasketSize', None, type=int,
84 info="default basket size for all branches in the output tree. "
85 "If not set (the default), no basket size is configured and ROOT's "
86 "default will be used.")
87

Member Function Documentation

◆ branchSortOrder()

python.OutputAnalysisConfig.OutputAnalysisConfig.branchSortOrder ( rule)
static

Definition at line 95 of file OutputAnalysisConfig.py.

95 def branchSortOrder (rule):
96 return rule.split('->')[1].strip()
97

◆ createOutputAlgs()

python.OutputAnalysisConfig.OutputAnalysisConfig.createOutputAlgs ( self,
config,
name,
vars,
nonContainers )
A helper function to create output algorithm

Definition at line 98 of file OutputAnalysisConfig.py.

98 def createOutputAlgs (self, config, name, vars, nonContainers):
99 """A helper function to create output algorithm"""
100 alg = config.createAlgorithm('CP::AsgxAODNTupleMakerAlg', name)
101 alg.TreeName = self.treeName
102 alg.RootStreamName = self.streamName
103 alg.NonContainers = list(nonContainers)
104 branchList = list(vars)
105 branchList.sort(key=self.branchSortOrder)
106 branchList_nosys = [branch for branch in branchList if "%SYS%" not in branch]
107 branchList_sys = [branch for branch in branchList if "%SYS%" in branch]
108 alg.Branches = branchList_nosys + branchList_sys
109 if self.defaultBasketSize is not None:
110 alg.DefaultBasketSize = self.defaultBasketSize
111 return alg
112

◆ createSelectionFlagBranches()

python.OutputAnalysisConfig.OutputAnalysisConfig.createSelectionFlagBranches ( self,
config,
containers )
For each container and for each selection, create a single pass variable in output NTuple,
which aggregates all the selections flag of the given selection. For example, this can include
pT, eta selections, some object ID selection, overlap removal, etc.
The goal is to have only one flag per object and working point in the output NTuple.

Definition at line 355 of file OutputAnalysisConfig.py.

355 def createSelectionFlagBranches(self, config, containers):
356 """
357 For each container and for each selection, create a single pass variable in output NTuple,
358 which aggregates all the selections flag of the given selection. For example, this can include
359 pT, eta selections, some object ID selection, overlap removal, etc.
360 The goal is to have only one flag per object and working point in the output NTuple.
361 """
362 originalContainersSeen = []
363 for prefix in containers.keys() :
364 outputContainerName = containers[prefix]
365 containerName = config.getOutputContainerOrigin(outputContainerName)
366 if containerName in originalContainersSeen:
367 continue
368 else:
369 originalContainersSeen.append(containerName)
370
371 # EventInfo is one obvious example of a container that has no object selections
372 if containerName == 'EventInfo':
373 continue
374
375 # Get the selection names, except the systematic-dependent version of the FTAG
376 # selection flag, as it's already saved as a systematic-independent output branch
377 selectionNames = config.getSelectionNames(containerName, excludeFrom={'ftag'})
378 for selectionName in selectionNames:
379 # skip default selection
380 if selectionName == '':
381 continue
382 # skip selection coming from the Thinning block
383 if self.skipRedundantSelectionFlags and "outputSelect" in selectionName:
384 continue
385 self.makeSelectionSummaryAlg(config, containerName, selectionName)
386

◆ instanceName()

python.OutputAnalysisConfig.OutputAnalysisConfig.instanceName ( self)
Return the instance name for this block

Definition at line 88 of file OutputAnalysisConfig.py.

88 def instanceName (self) :
89 """Return the instance name for this block"""
90 if self.postfix is not None and self.postfix != '':
91 return self.postfix
92 return self.treeName
93

◆ makeAlgs()

python.OutputAnalysisConfig.OutputAnalysisConfig.makeAlgs ( self,
config )

Definition at line 113 of file OutputAnalysisConfig.py.

113 def makeAlgs (self, config) :
114
115 log = logging.getLogger('OutputAnalysisConfig')
116
117 containers = dict(self.containers)
118 branchVars = set(self.vars)
119 metVars = set(self.metVars)
120 truthMetVars = set(self.truthMetVars)
121 nonContainers = list(self.nonContainers)
122 commands = list(self.commands)
123
124 # check for overlaps between containers and containersFullMET
125 overlapping_keys = set(containers.keys()).intersection(self.containersFullMET.keys())
126 if overlapping_keys:
127 # convert the set of overlapping keys to a list of strings for the message (represents the empty string too!)
128 keys_message = [repr(key) for key in overlapping_keys]
129 raise KeyError(f"containersFullMET would overwrite the following container keys: {', '.join(keys_message)}")
130 # move items in self.containersFullMET to containers
131 containers.update(self.containersFullMET)
132
133 # merge the MC-specific branches and containers into the main list/dictionary only if we are not running on data
134 if config.dataType() is not DataType.Data:
135 branchVars |= set(self.varsOnlyForMC)
136
137 # protect 'containers' against being overwritten
138 # find overlapping keys
139 overlapping_keys = set(containers.keys()).intersection(self.containersOnlyForMC.keys())
140 if overlapping_keys:
141 # convert the set of overlapping keys to a list of strings for the message (represents the empty string too!)
142 keys_message = [repr(key) for key in overlapping_keys]
143 raise KeyError(f"containersOnlyForMC would overwrite the following container keys: {', '.join(keys_message)}")
144
145 # move items in self.containersOnlyForMC to containers
146 containers.update(self.containersOnlyForMC)
147
148 # now filter the containers depending on DSIDs
149 if self.containersOnlyForDSIDs:
150 for container, dsid_filters in self.containersOnlyForDSIDs.items():
151 if container not in containers:
152 log.warning("Skipping unrecognised container prefix '%s' for DSID-filtering in OutputAnalysisConfig...", container)
153 continue
154 if not filter_dsids (dsid_filters, config):
155 # if current DSID is not allowed for this container, remove it
156 log.info("Skipping container prefix '%s' due to DSID filtering...", container)
157 # filter branches for validated containers
158 for var in set(branchVars): # make a copy of the list to avoid modifying it while iterating
159 var_container = var.split('.')[0].replace('_NOSYS', '').replace('_%SYS%', '')
160 if var_container == containers[container]:
161 branchVars.remove(var)
162 log.info("Skipping branch definition '%s' for excluded container %s...", var, var_container)
163 # filter branches for MET variables
164 for var in set(metVars): # make a copy of the list to avoid modifying it while iterating
165 var_container = var.split('.')[0].replace('_NOSYS', '').replace('_%SYS%', '')
166 if var_container == containers[container]:
167 metVars.remove(var)
168 log.info("Skipping MET branch definition '%s' for excluded container %s...", var, var_container)
169 # filter branches for truth MET variables
170 for var in set(truthMetVars): # make a copy of the list to avoid modifying it while iterating
171 var_container = var.split('.')[0].replace('_NOSYS', '').replace('_%SYS%', '')
172 if var_container == containers[container]:
173 truthMetVars.remove(var)
174 log.info("Skipping truth MET branch definition '%s' for excluded container %s...", var, var_container)
175 # remove the container from the list at the end
176 containers.pop (container)
177 # Filter individual variables depending on DSIDs
178 if self.varsOnlyForDSIDs:
179 for var_pattern, dsid_filters in self.varsOnlyForDSIDs.items():
180 if not filter_dsids(dsid_filters, config):
181 # Remove matching variables from user-defined vars collections
182 for var in set(branchVars):
183 if var_pattern in var:
184 branchVars.remove(var)
185 log.info("Skipping branch definition '%s' due to variable-level DSID filtering...", var)
186 for var in set(metVars):
187 if var_pattern in var:
188 metVars.remove(var)
189 log.info("Skipping MET branch definition '%s' due to variable-level DSID filtering...", var)
190 for var in set(truthMetVars):
191 if var_pattern in var:
192 truthMetVars.remove(var)
193 log.info("Skipping truth MET branch definition '%s' due to variable-level DSID filtering...", var)
194
195
196 for prefix, container in containers.items():
197 origName = config.getOutputContainerOrigin(container)
198 if config.getContainerMeta(origName, "nonContainer", False):
199 if origName not in nonContainers:
200 nonContainers.append(origName)
201
202 if self.storeSelectionFlags:
203 self.createSelectionFlagBranches(config, containers)
204
205 outputConfigs = {}
206 for prefix in containers.keys() :
207 containerName = containers[prefix]
208 outputDict = config.getOutputVars (containerName)
209 for outputName in outputDict :
210 outputConfig = copy.deepcopy (outputDict[outputName])
211 outputConfig.outputContainerName = config.readName(containerName)
212 outputConfig.prefix = prefix
213 # if the container is a MET container with all terms, we
214 # also need to write out the name of each MET term
215 if prefix in self.containersFullMET and outputConfig.variableName == 'name':
216 outputConfig.enabled = True
217 outputConfigs[prefix + outputName] = outputConfig
218
219 # check for DSID-specific commands
220 for dsid, dsid_commands in self.commandsOnlyForDSIDs.items():
221 if filter_dsids([dsid], config):
222 commands += dsid_commands
223
224 outputConfigsRename = {}
225 for command in commands :
226 words = command.split (' ')
227 if len (words) == 0 :
228 raise ValueError ('received empty command for "commands" option')
229 optional = words[0] == 'optional'
230 if optional :
231 words = words[1:] # remove the 'optional' keyword
232 if words[0] == 'enable' :
233 if len (words) != 2 :
234 raise ValueError ('enable takes exactly one argument: ' + command)
235 used = False
236 for name in outputConfigs :
237 if re.match (words[1], name) :
238 outputConfigs[name].enabled = True
239 used = True
240 if not used and not optional and config.dataType() is not DataType.Data:
241 raise KeyError ('unknown branch pattern for enable: ' + words[1])
242 elif words[0] == 'disable' :
243 if len (words) != 2 :
244 raise ValueError ('disable takes exactly one argument: ' + command)
245 used = False
246 for name in outputConfigs :
247 if re.match (words[1], name) :
248 outputConfigs[name].enabled = False
249 used = True
250 if not used and not optional and config.dataType() is not DataType.Data:
251 raise KeyError ('unknown branch pattern for disable: ' + words[1])
252 elif words[0] == 'rename' :
253 if len (words) != 3 :
254 raise ValueError ('rename takes exactly two arguments: ' + command)
255 used = False
256 for name in outputConfigs :
257 if re.match (words[1], name) :
258 new_name = re.sub (words[1], words[2], name)
259 outputConfigsRename[new_name] = copy.deepcopy(outputConfigs[name])
260 outputConfigs[name].enabled = False
261 used = True
262 if not used and not optional and config.dataType() is not DataType.Data:
263 raise KeyError ('unknown branch pattern for rename: ' + words[1])
264 else :
265 raise KeyError ('unknown command for "commands" option: ' + words[0])
266
267 # update the outputConfigs with renamed branches
268 outputConfigs.update(outputConfigsRename)
269
270 autoVars = set()
271 autoMetVars = set()
272 autoTruthMetVars = set()
273 for outputName, outputConfig in outputConfigs.items():
274 if outputConfig.enabled :
275 if config.isMetContainer (outputConfig.origContainerName) and outputConfig.prefix not in self.containersFullMET:
276 if "Truth" in outputConfig.origContainerName:
277 myVars = autoTruthMetVars
278 else:
279 myVars = autoMetVars
280 else :
281 myVars = autoVars
282 if outputConfig.noSys :
283 outputConfig.outputContainerName = outputConfig.outputContainerName.replace ('%SYS%', 'NOSYS')
284 outputConfig.variableName = outputConfig.variableName.replace ('%SYS%', 'NOSYS')
285 if self.alwaysAddNosys :
286 outputName += "_NOSYS"
287 else :
288 outputName += '_%SYS%'
289 branchDecl = f"{outputConfig.outputContainerName}.{outputConfig.variableName} -> {outputName}"
290 if outputConfig.auxType is not None :
291 branchDecl += f" type={outputConfig.auxType}"
292 if config.isMetContainer (outputConfig.origContainerName) and outputConfig.prefix not in self.containersFullMET:
293 if "Truth" in outputConfig.origContainerName:
294 branchDecl += f" metTerm={self.truthMetTermName}"
295 else:
296 branchDecl += f" metTerm={self.metTermName}"
297 myVars.add(branchDecl)
298
299 # Unified branch collection for all output formats
300 allBranches = set()
301 allBranches |= branchVars
302 allBranches |= autoVars
303 # Add MET branches
304 userMetVars = set()
305 if metVars:
306 for var in metVars:
307 userMetVars.add(var + " metTerm=" + self.metTermName)
308 allBranches |= userMetVars
309 allBranches |= autoMetVars
310 # Add truth MET branches (for MC)
311 userTruthMetVars = set()
312 if config.dataType() is not DataType.Data:
313 if truthMetVars:
314 for var in truthMetVars:
315 userTruthMetVars.add(var + " metTerm=" + self.truthMetTermName)
316 allBranches |= userTruthMetVars
317 allBranches |= autoTruthMetVars
318
319 # Create the output algorithm based on outputFormat
320 if self.outputFormat == 'RNTuple':
321 alg = config.createAlgorithm('CP::RNtupleTreeMakerAlg', 'RNtupleMaker')
322 alg.TreeName = self.treeName
323 alg.RootStreamName = self.streamName
324 alg.OutputStreamName = self.streamName
325 alg.NonContainers = list(nonContainers)
326
327 branchList = list(allBranches)
328 branchList.sort(key=self.branchSortOrder)
329 alg.Branches = branchList
330
331 return
332
333 # Add an ntuple dumper algorithm:
334 treeMaker = config.createAlgorithm( 'CP::TreeMakerAlg', 'TreeMaker' )
335 treeMaker.TreeName = self.treeName
336 treeMaker.RootStreamName = self.streamName
337 # the auto-flush setting still needs to be figured out
338 #treeMaker.TreeAutoFlush = 0
339
340 if branchVars or autoVars:
341 self.createOutputAlgs(config, 'NTupleMaker', branchVars | autoVars, nonContainers)
342
343 if metVars or autoMetVars:
344 self.createOutputAlgs(config, 'MetNTupleMaker', userMetVars | autoMetVars, nonContainers)
345
346 if config.dataType() is not DataType.Data and (truthMetVars or autoTruthMetVars):
347 self.createOutputAlgs(config, 'TruthMetNTupleMaker', userTruthMetVars | autoTruthMetVars, nonContainers)
348
349 treeFiller = config.createAlgorithm( 'CP::TreeFillerAlg', 'TreeFiller' )
350 treeFiller.TreeName = self.treeName
351 treeFiller.RootStreamName = self.streamName
352
353
354
STL class.
std::vector< std::string > intersection(std::vector< std::string > &v1, std::vector< std::string > &v2)
std::string replace(std::string s, const std::string &s2, const std::string &s3)
Definition hcg.cxx:312

◆ makeSelectionSummaryAlg()

python.OutputAnalysisConfig.OutputAnalysisConfig.makeSelectionSummaryAlg ( self,
config,
containerName,
selectionName )
Schedule an algorithm to pick up all cut flags for a given selectionName.
The summary selection flag is written to output as selectionFlagPrefix_selectionName.

Definition at line 387 of file OutputAnalysisConfig.py.

387 def makeSelectionSummaryAlg(self, config, containerName, selectionName):
388 """
389 Schedule an algorithm to pick up all cut flags for a given selectionName.
390 The summary selection flag is written to output as selectionFlagPrefix_selectionName.
391 """
392 alg = config.createAlgorithm( 'CP::AsgSelectionAlg',
393 f'ObjectSelectionSummary_{containerName}_{selectionName}')
394 selectionDecoration = f'baselineSelection_{selectionName}_%SYS%'
395 alg.selectionDecoration = f'{selectionDecoration},as_char'
396 alg.particles = config.readName (containerName)
397 alg.preselection = config.getFullSelection (containerName, selectionName)
398 config.addOutputVar (containerName, selectionDecoration, self.selectionFlagPrefix + '_' + selectionName)

Member Data Documentation

◆ alwaysAddNosys

python.OutputAnalysisConfig.OutputAnalysisConfig.alwaysAddNosys

Definition at line 285 of file OutputAnalysisConfig.py.

◆ containersFullMET

python.OutputAnalysisConfig.OutputAnalysisConfig.containersFullMET

Definition at line 131 of file OutputAnalysisConfig.py.

◆ containersOnlyForDSIDs

python.OutputAnalysisConfig.OutputAnalysisConfig.containersOnlyForDSIDs

Definition at line 149 of file OutputAnalysisConfig.py.

◆ containersOnlyForMC

python.OutputAnalysisConfig.OutputAnalysisConfig.containersOnlyForMC

Definition at line 146 of file OutputAnalysisConfig.py.

◆ metTermName

python.OutputAnalysisConfig.OutputAnalysisConfig.metTermName

Definition at line 307 of file OutputAnalysisConfig.py.

◆ outputFormat

str python.OutputAnalysisConfig.OutputAnalysisConfig.outputFormat = 'RNTuple':

Definition at line 320 of file OutputAnalysisConfig.py.

◆ storeSelectionFlags

python.OutputAnalysisConfig.OutputAnalysisConfig.storeSelectionFlags

Definition at line 202 of file OutputAnalysisConfig.py.

◆ truthMetTermName

python.OutputAnalysisConfig.OutputAnalysisConfig.truthMetTermName

Definition at line 315 of file OutputAnalysisConfig.py.

◆ varsOnlyForDSIDs

python.OutputAnalysisConfig.OutputAnalysisConfig.varsOnlyForDSIDs

Definition at line 178 of file OutputAnalysisConfig.py.


The documentation for this class was generated from the following file: