[TkAl][PixBary] Add "extract" job (->text files) from the output of a "single"

It runs extractBarycentre.py:
* input: ROOT files created by a "single" job
* output: text files that can be used to update the TWiki (csv, markdown)
  and a TeX table
* configurable styles (defaults are used for partition, quality etc.)

Updates in extractBarycentre.py:
* arg parsing outside main()
* support for writing directly to a file (instead of stdout)

[SQUASHME] iterate over all the single jobs

Author: amecca

Alignment/OfflineValidation/python/TkAlAllInOneTool/PixBary.py

@@ -13,6 +13,10 @@ def PixBary(config, validationDir, verbose=False):
     ##Start with single jobs
     jobType = "single"
 
+    ##Common paths
+    dir_base = os.path.join(validationDir, _validationName) # Base path for work directories (one for for each job, alignment and IOV)
+    out_base = os.path.join(config["LFS"], config["name"], _validationName)
+
     ##Check that a job is defined
     if not jobType in config["validations"][_validationName]:
         raise LookupError("No '%s' key word in config for %s" %(jobType, _validationName))
@@ -31,12 +35,9 @@ def PixBary(config, validationDir, verbose=False):
             for alignment in jobConfig["alignments"]:
                 alignmentConfig = config["alignments"][alignment]
 
-                ##Work directory for each IOV
-                workDir = os.path.join(validationDir, _validationName, jobType, jobName, alignment, IOV)
-
                 ##Write local config
                 local = {}
-                local["output"] = os.path.join(config["LFS"], config["name"], _validationName, jobType, alignment, jobName, IOV)
+                local["output"] = os.path.join(out_base, jobType, alignment, jobName, IOV)
                 local["alignment"] = copy.deepcopy(alignmentConfig)
                 local["alignment"]["label"] = alignment
                 local["validation"] = copy.deepcopy(jobConfig)
@@ -49,8 +50,9 @@ def PixBary(config, validationDir, verbose=False):
 
                 ##Write job info
                 job = {
+                    "label": jobName, # the name used in the config, so that it can be referenced lated
                     "name": "{}_{}_{}_{}_{}".format(_validationName, alignment, jobType, jobName, IOV),
-                    "dir": workDir,
+                    "dir": os.path.join(dir_base, jobType, jobName, alignment, IOV),
                     "exe": "cmsRun",
                     "cms-config": "{}/src/Alignment/OfflineValidation/python/TkAlAllInOneTool/PixelBaryCentreAnalyzer_cfg.py".format(os.environ["CMSSW_BASE"]),
                     "run-mode": "Condor",
@@ -60,6 +62,32 @@ def PixBary(config, validationDir, verbose=False):
 
                 jobs.append(job)
 
+    # Extract text from the ROOT files
+    jobType = "extract"
+
+    for jobName, jobConfig in config["validations"][_validationName][jobType].items():
+        for singleName in jobConfig.get("singles"):
+            # Search for the "single" job referenced by name
+            matchingSingleConfigs = [j for j in jobs if j.get("label", "") == singleName]
+
+            for singleConfig in matchingSingleConfigs:
+                IOV = singleConfig["config"]["validation"]["IOV"]  # <str>
+                alignment = singleConfig["config"]["alignment"]["label"] # <str>
+
+                job = {
+                    "name": "_".join([_validationName, jobType, jobName, singleName, alignment, IOV]),
+                    "dir": os.path.join(dir_base, jobType, jobName, singleName, alignment, IOV),
+                    "dependencies": [singleConfig["name"]],
+                    "exe": "extractBaryCentre.py",
+                    "config": {
+                        "input": os.path.join(singleConfig["config"]["output"], "PixelBaryCentre.root"),
+                        "output": os.path.join(out_base, jobType, jobName),
+                        "styles": jobConfig.get("styles", ["csv", "twiki"])
+                        },
+                    "flavour": "espresso" # So fast that anything else would not make sense
+                }
+                jobs.append(job)
+
     return jobs
 
 def get_IOVs(jobConfig):

Alignment/OfflineValidation/scripts/extractBaryCentre.py

@@ -2,6 +2,9 @@
 
 from argparse import ArgumentParser
 import logging
+import json
+import sys
+import os
 import ROOT
 
 class TFileContext(object):
@@ -12,7 +15,7 @@ def __enter__(self):
     def __exit__(self, type, value, traceback):
         self.tfile.Close()
 
-def display_results(t_data, style='twiki'):
+def display_results(t_data, style='twiki', out=sys.stdout):
     before = after = None
     if  (style == 'twiki'):
         header_fmt = ' | '.join( ['{:^6s}'] + ['{:^9s}' ] * (len(t_data.keys())-1) )
@@ -30,31 +33,18 @@ def display_results(t_data, style='twiki'):
     else:
         raise RuntimeError('Unknown style "%s" for table'%(style))
 
-    if(before is not None): print(before)
-    print( header_fmt.format(*t_data.keys()) )
+    if(before is not None): out.write(before+'\n')
+    out.write( header_fmt.format(*t_data.keys())+'\n' )
     for i, run in enumerate(t_data['run']):
-        print(row_fmt.format(run, *[t_data[k][i] for k in t_data.keys() if not k == 'run']))
-    if(after is not None): print(after)
+        out.write(row_fmt.format(run, *[t_data[k][i] for k in t_data.keys() if not k == 'run'])+'\n')
+    if(after is not None): out.write(after+'\n')
 
+    return out
 
-def main():
-    parser = ArgumentParser()
-    parser.add_argument('fname'            , metavar='FILE')
-    parser.add_argument('-p', '--partition', default='BPIX', help='Tracker partition (e.g. BPIX, FPIX, BPIXLYR1, etc.). Default: %(default)s')
-    parser.add_argument('-l', '--list-content' , action='store_true', dest='list_content', help='List the contents of file and exit')
-    parser.add_argument(      '--list-branches', action='store_true', help='List the branches of the tree and exit')
-    parser.add_argument('-t', '--type'     , default='barycentre', choices=('barycentre', 'beamspot'), type=str.lower, help='Default: %(default)s')
-    parser.add_argument(      '--label'    , default=None, help='Additional label that is appended to the folder name (i.e. PixelBaryCentreAnalyzer by default)')
-    parser.add_argument(      '--quality'  , action='store_true', help='Read results with the WithPixelQuality flag (default: %(default)s)')
-    parser.add_argument('-s', '--style'    , default='twiki', choices=('twiki', 'latex', 'csv'), type=str.lower, help='Table style for the output (default: %(default)s)')
-    parser.add_argument(      '--loglevel' , metavar='LEVEL', default='WARNING', help='Level for the python logging module. Can be either a mnemonic string like DEBUG, INFO or WARNING or an integer (lower means more verbose).')
 
-    args = parser.parse_args()
-    loglevel = args.loglevel.upper() if not args.loglevel.isdigit() else int(args.loglevel)
-    logging.basicConfig(format='%(levelname)s:%(module)s:%(funcName)s: %(message)s', level=loglevel)
+def main(args):
     logging.debug('args: %s', args)
 
-
     folder    = 'PixelBaryCentreAnalyzer' if not args.quality          else 'PixelBaryCentreAnalyzerWithPixelQuality'
     tree_name = 'PixelBarycentre'         if args.type == 'barycentre' else 'BeamSpot'
     if(args.label is not None):
@@ -85,7 +75,22 @@ def main():
         logging.info('Reading "%s"', tree_name)
         results = rdf.AsNumpy(columns)
 
-    display_results(results, style=args.style)
+    if(args.config is None):
+        display_results(results, style=args.style)
+    else:
+        # When this script is called from the validation framework, the output is
+        # written to a list of files, one for each style, with appropriate extensions
+        fname_base = '_'.join([args.type, args.partition] + (['quality'] if args.quality else []))
+        fname_path = os.path.join(args.config['output'], fname_base) # without the ext
+        for style in args.config["styles"]:
+            if  (style == 'twiki'): ext = 'txt'
+            elif(style == 'latex'): ext = 'tex'
+            else: ext = style
+            fname = '.'.join([fname_path, ext])
+            logging.info('output in "%s"', fname)
+
+            with open(fname, 'w') as fout:
+                display_results(results, style=style, out=fout)
 
     return 0
 
@@ -103,5 +108,36 @@ def list_branches(tree, folder_name=''):
     print('\n'.join(branches))
 
 
+def parse_args():
+    parser = ArgumentParser()
+    parser.add_argument('fname'            , metavar='FILE')
+    parser.add_argument('-p', '--partition', default='BPIX', help='Tracker partition (e.g. BPIX, FPIX, BPIXLYR1, etc.). Default: %(default)s')
+    parser.add_argument('-l', '--list-content' , action='store_true', dest='list_content', help='List the contents of file and exit')
+    parser.add_argument(      '--list-branches', action='store_true', help='List the branches of the tree and exit')
+    parser.add_argument('-t', '--type'     , default='barycentre', choices=('barycentre', 'beamspot'), type=str.lower, help='Default: %(default)s')
+    parser.add_argument(      '--label'    , default=None, help='Additional label that is appended to the folder name (i.e. PixelBaryCentreAnalyzer by default)')
+    parser.add_argument(      '--quality'  , action='store_true', help='Read results with the WithPixelQuality flag (default: %(default)s)')
+    parser.add_argument('-s', '--style'    , default='twiki', choices=('twiki', 'latex', 'csv'), type=str.lower, help='Table style for the output (default: %(default)s)')
+    parser.add_argument(      '--loglevel' , metavar='LEVEL', default='WARNING', help='Level for the python logging module. Can be either a mnemonic string like DEBUG, INFO or WARNING or an integer (lower means more verbose).')
+
+    args = parser.parse_args()
+
+    # If the script is called from the validation framework, the first argument
+    # will be a JSON file with the configuration, instead of a ROOT file
+    try:
+        with open(args.fname) as f:
+            config = json.load(f)
+        args.fname = config['input']
+        args.config = config
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        args.config = None
+
+    return args
+
+
 if __name__ == '__main__':
-    exit(main())
+    args = parse_args()
+    loglevel = args.loglevel.upper() if not args.loglevel.isdigit() else int(args.loglevel)
+    logging.basicConfig(format='%(levelname)s:%(module)s:%(funcName)s: %(message)s', level=loglevel)
+
+    exit(main(args))