Move the aggregated selection flag to output block and automatize

PhysicsAnalysis/Algorithms/AnalysisAlgorithmsConfig/python/ConfigAccumulator.py

@@ -577,6 +577,17 @@ class ConfigAccumulator :
         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
@@ -598,3 +609,16 @@ class ConfigAccumulator :
         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:
+            selectionNames.add(selection.name)
+        return selectionNames
\ No newline at end of file

PhysicsAnalysis/Algorithms/AsgAnalysisAlgorithms/python/OutputAnalysisConfig.py

@@ -16,6 +16,8 @@ class OutputAnalysisConfig (ConfigBlock):
         self.addOption ('containers', {}, type=None)
         self.addOption ('treeName', 'analysis', type=str)
         self.addOption ('metTermName', 'Final', type=str)
+        self.addOption ('storeSelectionFlags', True, type=bool)
+        self.addOption ('selectionFlagPrefix', 'passSelection', type=str)
         self.addOption ('systematicsHistogram', None , type=str)
         self.addOption ('commands', [], type=None,
                         info="a list of commands for branch selection/configuration")
@@ -23,6 +25,9 @@ class OutputAnalysisConfig (ConfigBlock):
 
     def makeAlgs (self, config) :
 
+        if self.storeSelectionFlags:
+            self.createSelectionFlagBranches(config)
+
         outputConfigs = {}
         for prefix in self.containers.keys() :
             containerName = self.containers[prefix]
@@ -107,3 +112,38 @@ class OutputAnalysisConfig (ConfigBlock):
         if self.systematicsHistogram is not None:
             sysDumper = config.createAlgorithm( 'CP::SysListDumperAlg', 'SystematicsPrinter' )
             sysDumper.histogramName = self.systematicsHistogram
+
+    def createSelectionFlagBranches(self, config):
+        """
+        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.
+        """
+        for prefix in self.containers.keys() :
+            outputContainerName = self.containers[prefix]
+            containerName = config.getOutputContainerOrigin(outputContainerName)
+
+            # EventInfo is one obvious example of a container that has no object selections
+            if containerName == 'EventInfo':
+                continue
+
+            selectionNames = config.getSelectionNames(containerName)
+            for selectionName in selectionNames:
+                # skip default selection
+                if selectionName == '':
+                    continue
+                self.makeSelectionSummaryAlg(config, containerName, selectionName)
+
+    def 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.
+        """
+        alg = config.createAlgorithm( 'CP::AsgSelectionAlg',
+                                      f'ObjectSelectionSummary_{containerName}_{selectionName}')
+        selectionFlagName = self.selectionFlagPrefix + '_' + selectionName
+        alg.selectionDecoration = selectionFlagName
+        alg.particles = config.readName (containerName)
+        alg.preselection = config.getFullSelection (containerName, selectionName)
+        config.addOutputVar (containerName, selectionFlagName, selectionFlagName)
\ No newline at end of file