src_JetMedianTree/macro/make_inp_lists.py

0001 #!/cvmfs/sphenix.sdcc.bnl.gov/gcc-12.1.0/opt/sphenix/core/bin/python3
0002
0003 '''
0004 Input:
0005     req: local directory full_lists/
0006     usage:
0007         make_inp_lists.py [num files total] [num files per list]
0008     output:
0009         new_dir: new directory named 'in_lists_[num files total]_by_[num files per list]'
0010         input files lists: [new_dir]/dst_*_[num files per list]_[0,1,...]
0011         new queue file: [queue_[num file total]_by_[num files per list].list]
0012 '''
0013
0014 from glob import glob
0015 from sys import argv
0016 import os
0017 from os import path
0018
0019
0020 def main():
0021     lines_per_output = 1 # lines per output
0022     try:
0023         lines_per_output = argv[1]
0024     except:
0025         pass
0026
0027     n_files = -1
0028     try:
0029         n_files = int(argv[2])
0030     except:
0031         pass
0032
0033     inp_dir = 'full_lists'
0034     try:
0035         inp_dir = argv[3]
0036     except:
0037         pass
0038
0039     # figure out the output directory tag
0040     odir_tag = "jobs"
0041     print (f"inp_dir: {inp_dir}")
0042     if not inp_dir == "full_lists":
0043         print (f"odir_tag: {odir_tag} 0")
0044         if "_lists" in inp_dir:
0045             odir_tag = inp_dir[:inp_dir.find("_lists")]
0046             print (f"odir_tag: {odir_tag} 1")
0047         else:
0048             odir_tag = "jobs"
0049             print (f"odir_tag: {odir_tag} 2")
0050
0051
0052     if n_files != -1:
0053         odir = f'{odir_tag}_{lines_per_output}x{n_files}'
0054     else:
0055         odir = f'{odir_tag}_{lines_per_output}xAll'
0056
0057     print (f"odir: {odir}")
0058
0059     if path.isdir(odir):
0060         print(f'replacing contents of {odir}');
0061         for f in glob(f'{odir}/*'):
0062             os.remove(f)
0063     else:
0064         os.mkdir(odir)
0065
0066     queue = open(f'{odir}/queue.list','w')
0067
0068     # determine which input files will be divided
0069     in_files = []
0070     in_tags  = []
0071     n_file = -1
0072     if (len(glob(f'{inp_dir}/dst_*.list')) == 0):
0073         print(f"fatal error: no lists in {inp_dir}")
0074         exit()
0075
0076     for file in glob(f'{inp_dir}/dst_*.list'):
0077         in_files .append(open(file,'r'))
0078         in_tags .append(file.split('/')[-1][:-5])
0079
0080     try:
0081         while True:
0082             n_file += 1
0083             if n_files != -1 and n_file == n_files:
0084                 break
0085             o_files = []
0086             # queue.write(f'out_{n_file}.root, ')
0087             first = True
0088             for tag in in_tags:
0089                 name = f'{tag}_{n_file}.list'
0090                 if first:
0091                     queue.write(f'{name}')
0092                     first = False
0093                 else:
0094                     queue.write(f', {name}')
0095                 o_files.append(open(f'{odir}/{name}','w'))
0096             queue.write('\n')
0097             for n in range (int(lines_per_output)):
0098                 for i_file, o_file in zip(in_files, o_files):
0099                     line = i_file.readline()
0100                     if not line:
0101                         for file in o_files:
0102                             o_file.close()
0103                         raise StopIteration
0104                     o_file.write(line)
0105             for file in o_files:
0106                 o_file.close()
0107     except StopIteration:
0108         pass
0109     if n_files == -1:
0110         with open(f'{odir}/last_list_is_{n_file}','w') as fout:
0111             fout.write(f'There are {n_file+2} lists')
0112
0113     queue.close()
0114
0115
0116
0117
0118 if __name__ == '__main__':
0119     main()