outputtxt2json.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-,

# __future__ imports must occur at beginning of file
# redirect python output using the newer print function with file description
#   print(string, f=fd)
from __future__ import print_function

# import the rest of the stuff
import argparse
import os
import sys
import csv
import collections
import json

def get_scaleFactor(weights, did):
  weight = weights.get(did, None)
  if weight is None:
    return 1.0
  scaleFactor = 1.0
  cutflow = weight.get('num events')
  if cutflow == 0:
    raise ValueError('Num events = 0!')
  scaleFactor /= cutflow
  scaleFactor *= weight.get('cross section', 1.0)
  scaleFactor *= weight.get('filter efficiency', 1.0)
  scaleFactor *= weight.get('k-factor', 1.0)
  return scaleFactor

if __name__ == "__main__":

  # if we want multiple custom formatters, use inheriting
  class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter):
    pass

  parser = argparse.ArgumentParser(description='Convert from the txt outputs to an output json combining all information',
                                   usage='\033[93m%(prog)s\033[0m files [options]',
                                   formatter_class=lambda prog: CustomFormatter(prog, max_help_position=30))
  parser.add_argument('files', type=str, nargs='+', help='Files to Convert')
  parser.add_argument('--analysis', type=str, required=True, help='Name of the analysis to strip off')
  parser.add_argument('--weights', metavar='weights.json', type=str, help='Weights file to weight all weighted and errs by')

  args = parser.parse_args()

  yields = collections.defaultdict(lambda: collections.defaultdict(lambda: collections.defaultdict(float)))
  weights = {}
  if args.weights:
    weights = json.load(file(args.weights))

  header_map = {"events": "raw", "acceptance": "weighted", "err": "err"}

  for fname in args.files:
    did = os.path.splitext(os.path.basename(fname))[0]
    sf = get_scaleFactor(weights, did)
    print("Reading in DID#{0:s} with SF {1:20.10f}".format(did, sf))
    with open(fname, 'r') as csvfile:
      reader = csv.reader(csvfile)
      headers = next(reader)
      headers[0] = None
      for row in reader:
        region = None
        for h, v in zip(headers, row):
          #if region in ['All']: continue
          if h is None:
            region = v.replace('{0:s}__'.format(args.analysis),'')
            print(h,'|', v, '|', region)
          else:
            val = float(v)
            if header_map[h] in ['weighted', 'err']:
              val*= sf
            yields[region][did][header_map[h]] += val

  with open('truth_dids.json', 'w+') as outfile:
    json.dump(yields, outfile, sort_keys=True, indent=4)