# Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
import AnaAlgorithm.DualUseConfig as DualUseConfig
from AthenaConfiguration.Enums import LHCPeriod, FlagEnum
import re
def mapUserName (name, *, noSysSuffix) :
"""map an internal name to a name for systematics data handles
Right now this just means appending a _%SYS% to the name."""
if not noSysSuffix :
return name + "_%SYS%"
else :
return name
class DataType(FlagEnum):
"""holds the various data types as an enum"""
Data = 'data'
FullSim = 'fullsim'
FastSim = 'fastsim'
class SelectionConfig :
"""all the data for a given selection that has been registered
the bits argument is for backward compatibility, does nothing, and will be
removed in the future."""
def __init__ (self, selectionName, decoration,
*, bits=0, preselection=None, comesFrom = '',
writeToOutput=True) :
self.name = selectionName
self.decoration = decoration
if preselection is not None :
self.preselection = preselection
else :
self.preselection = (selectionName == '')
self.comesFrom = comesFrom
self.writeToOutput = writeToOutput
class OutputConfig :
"""all the data for a given variables in the output that has been registered"""
def __init__ (self, origContainerName, variableName,
*, noSys, enabled) :
self.origContainerName = origContainerName
self.outputContainerName = None
self.variableName = variableName
self.noSys = noSys
self.enabled = enabled
class ContainerConfig :
"""all the auto-generated meta-configuration data for a single container
This tracks the naming of all temporary containers, as well as all the
selection decorations."""
def __init__ (self, name, sourceName, *, originalName = None, noSysSuffix) :
self.name = name
self.sourceName = sourceName
self.originalName = originalName
self.noSysSuffix = noSysSuffix
self.index = 0
self.maxIndex = None
self.viewIndex = 1
self.isMet = False
self.selections = []
self.outputs = {}
def currentName (self) :
if self.index == 0 :
if self.sourceName is None :
raise Exception ("should not get here, reading container name before created: " + self.name)
return self.sourceName
if self.maxIndex and self.index == self.maxIndex :
return mapUserName(self.name, noSysSuffix = self.noSysSuffix)
return mapUserName(self.name + "_STEP" + str(self.index), noSysSuffix = self.noSysSuffix)
def nextPass (self) :
self.maxIndex = self.index
self.index = 0
self.viewIndex = 1
self.selections = []
self.outputs = {}
class ConfigAccumulator :
"""a class that accumulates a configuration from blocks into an
algorithm sequence
This is used as argument to the ConfigurationBlock methods, which
need to be called in the correct order. This class will track all
meta-information that needs to be communicated between blocks
during configuration, and also add the created algorithms to the
sequence.
Use/access of containers in the event store is handled via
references that this class hands out. This happens in a separate
step before the algorithms are created, as the naming of
containers will depend on where in the chain the container is
used.
"""
def __init__ (self, algSeq, dataType=None, isPhyslite=False, geometry=None, dsid=0,
campaign=None, runNumber=None, autoconfigFromFlags=None, noSysSuffix=False,
noSystematics=None, dataYear=0):
self._autoconfigFlags = autoconfigFromFlags
if autoconfigFromFlags is not None:
if autoconfigFromFlags.Input.isMC:
if autoconfigFromFlags.Sim.ISF.Simulator.usesFastCaloSim():
dataType = DataType.FastSim
else:
dataType = DataType.FullSim
else:
dataType = DataType.Data
isPhyslite = 'StreamDAOD_PHYSLITE' in autoconfigFromFlags.Input.ProcessingTags
if geometry is None:
geometry = autoconfigFromFlags.GeoModel.Run
if dsid == 0 and dataType is not DataType.Data:
dsid = autoconfigFromFlags.Input.MCChannelNumber
if campaign is None:
campaign = autoconfigFromFlags.Input.MCCampaign
if runNumber is None:
runNumber = int(autoconfigFromFlags.Input.RunNumbers[0])
if dataYear == 0:
dataYear = autoconfigFromFlags.Input.DataYear
generatorInfo = autoconfigFromFlags.Input.GeneratorsInfo
else:
# legacy mappings of string arguments
if isinstance(dataType, str):
if dataType == 'mc':
dataType = DataType.FullSim
elif dataType == 'afii':
dataType = DataType.FastSim
else:
dataType = DataType(dataType)
generatorInfo = None
if runNumber is None:
runNumber = 284500
# allow possible string argument for `geometry` and convert it to enum
geometry = LHCPeriod(geometry)
if geometry is LHCPeriod.Run1:
raise ValueError ("invalid Run geometry: %s" % geometry.value)
# store also the data year for data
self._dataType = dataType
self._isPhyslite = isPhyslite
self._geometry = geometry
self._dsid = dsid
self._campaign = campaign
self._runNumber = runNumber
self._dataYear = dataYear
self._generatorInfo = generatorInfo
self._algSeq = algSeq
self._noSystematics = noSystematics
self._noSysSuffix = noSysSuffix
self._containerConfig = {}
self._outputContainers = {}
self._pass = 0
self._algorithms = {}
self._currentAlg = None
self._selectionNameExpr = re.compile ('[A-Za-z_][A-Za-z_0-9]+')
self.setSourceName ('EventInfo', 'EventInfo')
self._eventcutflow = {}
self._muonCalibMode = -1
# If we are in an Athena environment with ComponentAccumulator configuration
# then the AlgSequence, which is Gaudi.AthSequencer, does not support '+=',
# and we in any case want to produce an output ComponentAccumulator
self.CA = None
if DualUseConfig.useComponentAccumulator:
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
self.CA = ComponentAccumulator()
# if we have a component accumulator the user is not required to pass
# in a sequence, but if they do let's add it
if algSeq :
self.CA.addSequence(algSeq)
def noSystematics (self) :
"""noSystematics flag used by CommonServices block"""
return self._noSystematics
def autoconfigFlags (self) :
"""auto configuration flags"""
return self._autoconfigFlags
def dataType (self) :
"""the data type we run on (data, fullsim, fastsim)"""
return self._dataType
def isPhyslite (self) :
"""whether we run on PHYSLITE"""
return self._isPhyslite
def geometry (self) :
"""the LHC Run period we run on"""
return self._geometry
def dsid(self) :
"""the mcChannelNumber or DSID of the sample we run on"""
return self._dsid
def campaign(self) :
"""the MC campaign we run on"""
return self._campaign
def runNumber(self) :
"""the MC runNumber"""
return self._runNumber
def dataYear(self) :
"""for data, the corresponding year; for MC, zero"""
return self._dataYear
def generatorInfo(self) :
"""the dictionary of MC generators and their versions for the sample we run on"""
return self._generatorInfo
def createAlgorithm (self, type, name, reentrant=False) :
"""create a new algorithm and register it as the current algorithm"""
if self._pass == 0 :
if name in self._algorithms :
raise Exception ('duplicate algorithms: ' + name)
if reentrant:
alg = DualUseConfig.createReentrantAlgorithm (type, name)
else:
alg = DualUseConfig.createAlgorithm (type, name)
if DualUseConfig.useComponentAccumulator:
if self._algSeq :
self.CA.addEventAlgo(alg,self._algSeq.name)
else :
self.CA.addEventAlgo(alg)
else:
self._algSeq += alg
self._algorithms[name] = alg
self._currentAlg = alg
return alg
else :
if name not in self._algorithms :
raise Exception ('unknown algorithm requested: ' + name)
self._currentAlg = self._algorithms[name]
if self.CA and self._currentAlg != self.CA.getEventAlgo(name) :
raise Exception ('change to algorithm object: ' + name)
return self._algorithms[name]
def createService (self, type, name) :
'''create a new service and register it as the "current algorithm"'''
if self._pass == 0 :
if name in self._algorithms :
raise Exception ('duplicate service: ' + name)
service = DualUseConfig.createService (type, name)
# Avoid importing AthenaCommon.AppMgr in a CA Athena job
# as it modifies Gaudi behaviour
if DualUseConfig.isAthena:
if DualUseConfig.useComponentAccumulator:
self.CA.addService(service)
else:
# We're not, so let's remember this as a "normal" algorithm:
self._algSeq += service
self._algorithms[name] = service
self._currentAlg = service
return service
else :
if name not in self._algorithms :
raise Exception ('unknown service requested: ' + name)
self._currentAlg = self._algorithms[name]
return self._algorithms[name]
def createPublicTool (self, type, name) :
'''create a new public tool and register it as the "current algorithm"'''
if self._pass == 0 :
if name in self._algorithms :
raise Exception ('duplicate public tool: ' + name)
tool = DualUseConfig.createPublicTool (type, name)
# Avoid importing AthenaCommon.AppMgr in a CA Athena job
# as it modifies Gaudi behaviour
if DualUseConfig.isAthena:
if DualUseConfig.useComponentAccumulator:
self.CA.addPublicTool(tool)
else:
# We're not, so let's remember this as a "normal" algorithm:
self._algSeq += tool
self._algorithms[name] = tool
self._currentAlg = tool
return tool
else :
if name not in self._algorithms :
raise Exception ('unknown public tool requested: ' + name)
self._currentAlg = self._algorithms[name]
return self._algorithms[name]
def addPrivateTool (self, propertyName, toolType) :
"""add a private tool to the current algorithm"""
if self._pass == 0 :
DualUseConfig.addPrivateTool (self._currentAlg, propertyName, toolType)
def setSourceName (self, containerName, sourceName,
*, originalName = None) :
"""set the (default) name of the source/original container
This is essentially meant to allow using e.g. the muon
configuration and the user not having to manually specify that
they want to use the Muons/AnalysisMuons container from the
input file.
In addition it allows to set the original name of the
container (which may be different from the source name), which
is mostly/exclusively used for jet containers, so that
subsequent configurations know which jet container they
operate on.
"""
if containerName not in self._containerConfig :
self._containerConfig[containerName] = ContainerConfig (containerName, sourceName, noSysSuffix = self._noSysSuffix, originalName = originalName)
def writeName (self, containerName, *, isMet=None) :
"""register that the given container will be made and return
its name"""
if containerName not in self._containerConfig :
self._containerConfig[containerName] = ContainerConfig (containerName, sourceName = None, noSysSuffix = self._noSysSuffix)
if self._containerConfig[containerName].sourceName is not None :
raise Exception ("trying to write container configured for input: " + containerName)
if self._containerConfig[containerName].index != 0 :
raise Exception ("trying to write container twice: " + containerName)
self._containerConfig[containerName].index += 1
if isMet is not None :
self._containerConfig[containerName].isMet = isMet
return self._containerConfig[containerName].currentName()
def readName (self, containerName) :
"""get the name of the "current copy" of the given container
As extra copies get created during processing this will track
the correct name of the current copy. Optionally one can pass
in the name of the container before the first copy.
"""
if containerName not in self._containerConfig :
raise Exception ("no source container for: " + containerName)
return self._containerConfig[containerName].currentName()
def copyName (self, containerName) :
"""register that a copy of the container will be made and return
its name"""
if containerName not in self._containerConfig :
raise Exception ("unknown container: " + containerName)
self._containerConfig[containerName].index += 1
return self._containerConfig[containerName].currentName()
def wantCopy (self, containerName) :
"""ask whether we want/need a copy of the container
This usually only happens if no copy of the container has been
made yet and the copy is needed to allow modifications, etc.
"""
if containerName not in self._containerConfig :
raise Exception ("no source container for: " + containerName)
return self._containerConfig[containerName].index == 0
def originalName (self, containerName) :
"""get the "original" name of the given container
This is mostly/exclusively used for jet containers, so that
subsequent configurations know which jet container they
operate on.
"""
if containerName not in self._containerConfig :
raise Exception ("container unknown: " + containerName)
result = self._containerConfig[containerName].originalName
if result is None :
raise Exception ("no original name for: " + containerName)
return result
def isMetContainer (self, containerName) :
"""whether the given container is registered as a MET container
This is mostly/exclusively used for determining whether to
write out the whole container or just a single MET term.
"""
if containerName not in self._containerConfig :
raise Exception ("container unknown: " + containerName)
return self._containerConfig[containerName].isMet
def readNameAndSelection (self, containerName, *, excludeFrom = None) :
"""get the name of the "current copy" of the given container, and the
selection string
This is mostly meant for MET and OR for whom the actual object
selection is relevant, and which as such allow to pass in the
working point as "ObjectName.WorkingPoint".
"""
split = containerName.split (".")
if len(split) == 1 :
objectName = split[0]
selectionName = ''
elif len(split) == 2 :
objectName = split[0]
selectionName = split[1]
else :
raise Exception ('invalid object selection name: ' + containerName)
return self.readName (objectName), self.getFullSelection (objectName, selectionName, excludeFrom=excludeFrom)
def nextPass (self) :
"""switch to the next configuration pass
Configuration happens in two steps, with all the blocks processed
twice. This switches from the first to the second pass.
"""
if self._pass != 0 :
raise Exception ("already performed final pass")
for name in self._containerConfig :
self._containerConfig[name].nextPass ()
self._pass = 1
self._currentAlg = None
self._outputContainers = {}
def getPreselection (self, containerName, selectionName, *, asList = False) :
"""get the preselection string for the given selection on the given
container
"""
if selectionName != '' and not self._selectionNameExpr.fullmatch (selectionName) :
raise ValueError ('invalid selection name: ' + selectionName)
if containerName not in self._containerConfig :
return ""
config = self._containerConfig[containerName]
decorations = []
for selection in config.selections :
if (selection.name == '' or selection.name == selectionName) and \
selection.preselection :
decorations += [selection.decoration]
if asList :
return decorations
else :
return '&&'.join (decorations)
def getFullSelection (self, containerName, selectionName,
*, skipBase = False, excludeFrom = None) :
"""get the selection string for the given selection on the given
container
This can handle both individual selections or selection
expressions (e.g. `loose||tight`) with the later being
properly expanded. Either way the base selection (i.e. the
selection without a name) will always be applied on top.
containerName --- the container the selection is defined on
selectionName --- the name of the selection, or a selection
expression based on multiple named selections
skipBase --- will avoid the base selection, and should normally
not be used by the end-user.
excludeFrom --- a set of string names of selection sources to exclude
e.g. to exclude OR selections from MET
"""
if containerName not in self._containerConfig :
return ""
if excludeFrom is None :
excludeFrom = set()
elif not isinstance(excludeFrom, set) :
raise ValueError ('invalid excludeFrom argument (need set of strings): ' + str(excludeFrom))
# Check if this is actually a selection expression,
# e.g. `A||B` and if so translate it into a complex expression
# for the user. I'm not trying to do any complex syntax
# recognition, but instead just produce an expression that the
# C++ parser ought to be able to read.
if selectionName != '' and \
not self._selectionNameExpr.fullmatch (selectionName) :
result = ''
while selectionName != '' :
match = self._selectionNameExpr.match (selectionName)
if not match :
result += selectionName[0]
selectionName = selectionName[1:]
else :
subname = match.group(0)
subresult = self.getFullSelection (containerName, subname, skipBase = True, excludeFrom=excludeFrom)
if subresult != '' :
result += '(' + subresult + ')'
else :
result += 'true'
selectionName = selectionName[len(subname):]
subresult = self.getFullSelection (containerName, '', excludeFrom=excludeFrom)
if subresult != '' :
result = subresult + '&&(' + result + ')'
return result
config = self._containerConfig[containerName]
decorations = []
hasSelectionName = False
for selection in config.selections :
if ((selection.name == '' and not skipBase) or selection.name == selectionName) and (selection.comesFrom not in excludeFrom) :
decorations += [selection.decoration]
if selection.name == selectionName :
hasSelectionName = True
if not hasSelectionName and selectionName != '' :
raise KeyError ('invalid selection name: ' + containerName + '.' + selectionName)
return '&&'.join (decorations)
def getSelectionCutFlow (self, containerName, selectionName) :
"""get the individual selections as a list for producing the cutflow for
the given selection on the given container
This can only handle individual selections, not selection
expressions (e.g. `loose||tight`).
"""
if containerName not in self._containerConfig :
return []
# Check if this is actually a selection expression,
# e.g. `A||B` and if so translate it into a complex expression
# for the user. I'm not trying to do any complex syntax
# recognition, but instead just produce an expression that the
# C++ parser ought to be able to read.
if selectionName != '' and \
not self._selectionNameExpr.fullmatch (selectionName) :
raise ValueError ('not allowed to do cutflow on selection expression: ' + selectionName)
config = self._containerConfig[containerName]
decorations = []
for selection in config.selections :
if (selection.name == '' or selection.name == selectionName) :
decorations += [selection.decoration]
return decorations
def addEventCutFlow (self, selection, decorations) :
"""register a new event cutflow, adding it to the dictionary with key 'selection'
and value 'decorations', a list of decorated selections
"""
if self._pass == 0:
if selection in self._eventcutflow.keys():
raise ValueError ('the event cutflow dictionary already contains an entry ' + selection)
else:
self._eventcutflow[selection] = decorations
def getEventCutFlow (self, selection) :
"""get the list of decorated selections for an event cutflow, corresponding to
key 'selection'
"""
return self._eventcutflow[selection]
def addSelection (self, containerName, selectionName, decoration,
**kwargs) :
"""add another selection decoration to the selection of the given
name for the given container"""
if selectionName != '' and not self._selectionNameExpr.fullmatch (selectionName) :
raise ValueError ('invalid selection name: ' + selectionName)
if containerName not in self._containerConfig :
self._containerConfig[containerName] = ContainerConfig (containerName, containerName, noSysSuffix=self._noSysSuffix)
config = self._containerConfig[containerName]
selection = SelectionConfig (selectionName, decoration, **kwargs)
config.selections.append (selection)
def addOutputContainer (self, containerName, outputContainerName) :
"""register a copy of a container used in outputs"""
if containerName not in self._containerConfig :
raise KeyError ("container unknown: " + containerName)
if outputContainerName in self._outputContainers :
raise KeyError ("duplicate output container name: " + outputContainerName)
self._outputContainers[outputContainerName] = containerName
def getOutputContainerOrigin (self, outputContainerName) :
"""Get the name of the actual container, for which an output is registered"""
try:
return self._outputContainers[outputContainerName]
except KeyError:
try:
return self._containerConfig[outputContainerName].name
except KeyError:
raise KeyError ("output container unknown: " + outputContainerName)
def addOutputVar (self, containerName, variableName, outputName,
*, noSys=False, enabled=True) :
"""add an output variable for the given container to the output
"""
if containerName not in self._containerConfig :
raise KeyError ("container unknown: " + containerName)
baseConfig = self._containerConfig[containerName].outputs
if outputName in baseConfig :
raise KeyError ("duplicate output variable name: " + outputName)
config = OutputConfig (containerName, variableName, noSys=noSys, enabled=enabled)
baseConfig[outputName] = config
def getOutputVars (self, containerName) :
"""get the output variables for the given container"""
if containerName in self._outputContainers :
containerName = self._outputContainers[containerName]
if containerName not in self._containerConfig :
raise KeyError ("unknown container for output: " + containerName)
return self._containerConfig[containerName].outputs
def getSelectionNames (self, containerName) :
"""Retrieve set of unique selections defined for a given container"""
if containerName not in self._containerConfig :
return []
config = self._containerConfig[containerName]
# because cuts are registered individually, selection names can repeat themselves
# but we are interested in unique names only
selectionNames = set()
for selection in config.selections:
# skip flags which should be disabled in output
if selection.writeToOutput:
selectionNames.add(selection.name)
return selectionNames