Merge branch '21.0-update-trfArgClasses' into '21.0'

21.0 update trf arg classes

See merge request atlas/athena!11987

Former-commit-id: d0b0739cc87ac012f2100e85181f71b766874312

Tools/PyJobTransforms/python/trfArgClasses.py

@@ -21,7 +21,7 @@ msg = logging.getLogger(__name__)
 
 import PyJobTransforms.trfExceptions as trfExceptions
 
-from PyJobTransforms.trfFileUtils import athFileInterestingKeys, AthenaLiteFileInfo, NTUPEntries, HISTEntries, urlType, ROOTGetSize
+from PyJobTransforms.trfFileUtils import athFileInterestingKeys, AthenaLiteFileInfo, NTUPEntries, HISTEntries, PRWEntries, urlType, ROOTGetSize
 from PyJobTransforms.trfUtils import call, cliToKey
 from PyJobTransforms.trfExitCodes import trfExit as trfExit
 from PyJobTransforms.trfDecorators import timelimited
@@ -1578,14 +1578,26 @@ class argNTUPFile(argFile):
                                    'file_guid': self._generateGUID,
                                    'integrity': self._getIntegrity,
                                    })
-                
-        
+
+        if name and 'NTUP_PILEUP' in name:
+            self._metadataKeys.update({
+                                       'sumOfWeights': self._getNumberOfEvents,
+                                       })
+
     def _getNumberOfEvents(self, files):
         msg.debug('Retrieving event count for NTUP files {0}'.format(files))
         if self._treeNames is None:
-            msg.debug('treeNames is set to None - event count undefined for this NTUP')
             for fname in files:
-                self._fileMetadata[fname]['nentries'] = 'UNDEFINED'
+                # Attempt to treat this as a pileup reweighting file
+                myEntries = PRWEntries(fileName=fname)
+                if myEntries is not None:
+                    self._fileMetadata[fname]['nentries'] = myEntries
+                    if self.name and 'NTUP_PILEUP' in self.name:
+                        myEntries = PRWEntries(fileName=fname, integral=True)
+                        self._fileMetadata[fname]['sumOfWeights'] = myEntries
+                else:
+                    msg.debug('treeNames is set to None - event count undefined for this NTUP')
+                    self._fileMetadata[fname]['nentries'] = 'UNDEFINED'
         else:
             for fname in files:
                 try:
@@ -1718,6 +1730,70 @@ class argHepEvtAsciiFile(argFile):
                 msg.error('Event count for file {0} failed: {1!s}'.format(fname, e))
                 self._fileMetadata[fname]['nentries'] = None
                 
+## @brief LHE ASCII file 
+class argLHEFile(argFile):
+    def __init__(self, value=list(), io = 'output', type=None, splitter=',', runarg=True, multipleOK=None, name=None):
+        super(argLHEFile, self).__init__(value=value, io=io, type=type, splitter=splitter, runarg=runarg, multipleOK=multipleOK,
+                                           name=name)
+
+        self._metadataKeys.update({
+                'nentries': self._getNumberOfEvents,
+                'lheSumOfPosWeights': self._getWeightedEvents,
+                'lheSumOfNegWeights': 0,
+                })
+
+    def _getNumberOfEvents(self, files):
+        msg.debug('Retrieving event count for LHE file {0}'.format(files))
+        import tarfile
+        for fname in files:
+            # Attempt to treat this as a pileup reweighting file
+            try :
+                tar = tarfile.open(fname, "r:gz")
+                lhecount = 0
+                for untar in tar.getmembers():
+                    fileTXT = tar.extractfile(untar)
+                    if fileTXT is not None :
+                        lines = fileTXT.read()
+                        lhecount = lines.find('/event')
+
+                self._fileMetadata[fname]['nentries'] = lhecount
+            except :
+                msg.debug('Entries is set to None - event count undefined for this LHE')
+                self._fileMetadata[fname]['nentries'] = 'UNDEFINED'
+
+    def _getWeightedEvents(self, files):
+        msg.debug('Retrieving weight count for LHE file {0}'.format(files))
+        import tarfile
+        import re
+
+        for fname in files:
+            weightPos = 0
+            weightNeg = 0
+            try :
+                tar = tarfile.open(fname, "r:gz")
+                for untar in tar.getmembers():
+                    fileTXT = tar.extractfile(untar)
+                    next = False
+                    if fileTXT is not None :
+                        lines = fileTXT.readlines()
+                        for line in lines :
+                            if next :
+                                try :
+                                    w = float(re.sub(' +',' ',line).split(" ")[2])
+                                    if w > 0 : weightPos += w
+                                    else : weightNeg += abs(w)
+                                except :
+                                    pass
+                                next = False
+                            if "<event" in line :
+                                next = True
+
+                self._fileMetadata[fname]['lheSumOfPosWeights'] = weightPos
+                self._fileMetadata[fname]['lheSumOfNegWeights'] = weightNeg
+            except :
+                msg.debug('Entries is set to None - negative fraction count undefined for this LHE')
+                self._fileMetadata[fname]['lheSumOfPosWeights'] = 'UNDEFINED'
+                self._fileMetadata[fname]['lheSumOfNegWeights'] = 'UNDEFINED'
 
 ## @brief Base class for substep arguments
 #  @details Sets up a dictionary with {substep1: value1, substep2: value2, ...}

Tools/PyJobTransforms/python/trfArgs.py

@@ -110,8 +110,9 @@ def addAthenaArguments(parser, maxEventsDefaultSubstep='first', addValgrind=True
                         metavar='BOOL', group='Athena',
                         help='SharedWriter mode active')
     parser.add_argument('--eventService', type=trfArgClasses.argFactory(trfArgClasses.argBool, runarg=True),
-                         metavar='BOOL', group='Athena',
-                         help='Switch AthenaMP to the Event Service configuration')
+                        metavar='BOOL', group='Athena',
+                        help='Switch AthenaMP to the Event Service configuration')
+
     if addValgrind:
         addValgrindArguments(parser)
 
@@ -474,8 +475,8 @@ def addExtraDPDTypes(parser, pick=None, transform=None, multipleOK=False, NTUPMe
                 msg.debug('Adding DPD {0} ({1}, {2}, {3}, {4})'.format(dpd.name, dpd.type, dpd.substeps, dpd.treeNames, dpd.argclass))
                 # NTUPs are a bit special as they can take a treeName to count events
                 if issubclass(dpd.argclass, trfArgClasses.argNTUPFile):
-                    parser.add_argument('--output' + dpd.name + 'File', 
-                                        type=argFactory(dpd.argclass, multipleOK=multipleOK, type=dpd.type, treeNames=dpd.treeNames), 
+                    parser.add_argument('--output' + dpd.name + 'File',
+                                        type=argFactory(dpd.argclass, name=dpd.name.upper(), multipleOK=multipleOK, type=dpd.type, treeNames=dpd.treeNames), 
                                         group = 'Additional DPDs', metavar=dpd.name.upper(), 
                                         help=dpd.help if dpd.help else 'DPD output {0} file'.format(dpd.name))
                 else:

Tools/PyJobTransforms/python/trfFileUtils.py

@@ -212,6 +212,48 @@ def NTUPEntries(fileName, treeNames):
     return numberOfEntries
 
 
+## @brief Determines number of entries in PRW file
+#  @param fileName Path to the PRW file.
+#  @param integral Returns sum of weights if true
+#  @return 
+#  - Number of entries.
+#  - Sum of weights if integral is true.
+#  - @c None if the determination failed.
+#  @note Use the PyCmt forking decorator to ensure that ROOT is run completely within 
+#  a child process and will not 'pollute' the parent python process with unthread-safe
+#  bits of code (otherwise strange hangs are observed on subsequent uses of ROOT)
+@_decos.forking
+def PRWEntries(fileName, integral=False):
+
+    root = import_root()
+
+    fname = root.TFile.Open(fileName, 'READ')
+
+    if not (isinstance(fname, root.TFile) and fname.IsOpen()):
+        return None
+
+    rundir = None
+
+    for key in fname.GetListOfKeys():
+        if key.GetName()=='PileupReweighting':
+            rundir = fname.Get('PileupReweighting')
+            break
+        # Not PRW...
+
+    if rundir is None: return None
+
+    total = 0
+    for key in rundir.GetListOfKeys():
+        if 'pileup' in key.GetName():
+            msg.debug('Working on file '+fileName+' histo '+key.GetName())
+            if integral:
+                total += rundir.Get(key.GetName()).Integral()
+            else:
+                total += rundir.Get(key.GetName()).GetEntries()
+        # Was not one of our histograms
+    return total
+
+
 ## @brief Get the size of a file via ROOT's TFile
 #  @details Use TFile.Open to retrieve a ROOT filehandle, which will
 #  deal with all non-posix filesystems. Return the GetSize() value.