BTau/python/generateFlavCfromFlavB.py

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