Skip to content

Instantly share code, notes, and snippets.

@kratsg
Created February 25, 2023 13:23
Show Gist options
  • Save kratsg/675e2c46aecbf399bd4e6095c4738d72 to your computer and use it in GitHub Desktop.
Save kratsg/675e2c46aecbf399bd4e6095c4738d72 to your computer and use it in GitHub Desktop.
#!/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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment