BTau/python/generateFlavCfromFlavB.py

0001 #!/usr/bin/env python
0002
0003 from __future__ import print_function
0004 from __future__ import absolute_import
0005 import os
0006 import sys
0007 import itertools
0008 from . import dataLoader
0009 from . import checkBTagCalibrationConsistency as checker
0010
0011
0012 def generate_flav_c(loaded_data):
0013     flav_b_data = [e for e in loaded_data.entries if e.params.jetFlavor == 0]
0014     flav_b_data = sorted(flav_b_data, key=lambda e: e.params.operatingPoint)
0015     flav_b_data = sorted(flav_b_data, key=lambda e: e.params.measurementType)
0016     flav_b_data = sorted(flav_b_data, key=lambda e: e.params.etaMin)
0017     flav_b_data = sorted(flav_b_data, key=lambda e: e.params.ptMin)
0018     flav_b_data = sorted(flav_b_data, key=lambda e: e.params.discrMin)
0019     sys_groups = itertools.groupby(
0020         flav_b_data,
0021         key=lambda e: '%d, %s, %.02f, %.02f, %.02f' % (
0022             e.params.operatingPoint,
0023             e.params.measurementType,
0024             e.params.etaMin,
0025             e.params.ptMin,
0026             e.params.discrMin
0027         )
0028     )
0029
0030     def gen_entry_dict(groups):
0031         for _, grp in groups:
0032             grp = list(grp)
0033             entries_by_sys = dict((e.params.sysType, e) for e in grp)
0034             assert len(grp) == len(entries_by_sys)  # every sysType is unique
0035             yield entries_by_sys
0036     sys_dicts = gen_entry_dict(sys_groups)
0037
0038     def gen_flavb_csv_line(dicts):
0039         for d in dicts:
0040             central = d.pop('central')
0041             central.params.jetFlavor = 1
0042             yield central.makeCSVLine()
0043             for e in d.values():
0044                 e.params.jetFlavor = 1
0045                 e.formula = '2*(%s)-(%s)' % (e.formula, central.formula)
0046                 yield e.makeCSVLine()
0047     csv_lines = gen_flavb_csv_line(sys_dicts)
0048
0049     return list(l for l in csv_lines)
0050
0051
0052 def main():
0053     if len(sys.argv) < 3:
0054         print('Need input/output filenames as first/second arguments. Exit.')
0055         exit(-1)
0056     if os.path.exists(sys.argv[2]):
0057         print('Output file exists. Exit.')
0058         exit(-1)
0059
0060     print('\nChecking input file consistency...')
0061     loaders = dataLoader.get_data(sys.argv[1])
0062     checks = checker.run_check_data(loaders, True, True, False)
0063     for data in loaders:
0064         typ = data.meas_type
0065         if 1 in data.flavs:
0066             print('FLAV_C already present in input file for %s. Exit.' % typ)
0067             exit(-1)
0068     if not any(0 in data.flavs for data in loaders):
0069         print('FLAV_B not found in input file. Exit.')
0070         exit(-1)
0071
0072
0073     print('\nGenerating new csv content...')
0074     new_csv_data = list(itertools.chain.from_iterable(
0075         l
0076         for d in loaders
0077         for l in generate_flav_c(d)
0078     ))
0079
0080     with open(sys.argv[1]) as f:
0081         old_csv_data = f.readlines()
0082
0083     with open(sys.argv[2], 'w') as f:
0084         f.writelines(old_csv_data)
0085         f.write('\n')
0086         f.writelines(new_csv_data)
0087
0088     print('Done.')
0089
0090
0091 if __name__ == '__main__':
0092     main()