Corral/lists/make_lists.py

0001 #!/usr/bin/env python3
0002 # Make corral job lists from a segment_entries.txt (written by count_entries.C):
0003 #   one list = contiguous segments of ONE run, in segment order, ~TARGET events.
0004 #   Per run: nlists = max(1, round(run_events/TARGET)); cut the run into nlists contiguous groups
0005 #   of nearly equal event counts. Lists are numbered 0..N-1 across runs (run order, then segment
0006 #   order) = the SLURM array index. Writes <outdir>/list_NN.txt (one file path per line),
0007 #   <outdir>/lists_summary.txt, and <outdir>/manifest.txt.
0008 #   Usage: make_lists.py <outdir> <target_events_per_list> <collect_dir>
0009 import os, sys, collections, datetime
0010 outdir, TARGET, cdir = sys.argv[1], int(sys.argv[2]), sys.argv[3]
0011 rows   = [l.split() for l in open(os.path.join(outdir, "segment_entries.txt")) if l.strip()]
0012 bad    = [r for r in rows if int(r[2]) <= 0]
0013 if bad:
0014     print("unreadable/empty files (left out of the lists):"); print("\n".join(r[3] for r in bad))
0015 runs   = collections.OrderedDict()
0016 for run, seg, ne, path in rows:
0017     if int(ne) <= 0: continue
0018     runs.setdefault(int(run), []).append((int(seg), int(ne), path))
0019 for r in runs: runs[r].sort()
0020 ilist, summary = 0, []
0021 for r, segs in runs.items():
0022     tot    = sum(s[1] for s in segs)
0023     nl     = max(1, int(round(tot / TARGET)))
0024     cuts   = [tot * (k + 1) / nl for k in range(nl)]    # cumulative-event targets for the end of each list
0025     groups, cur, acc, k = [], [], 0, 0
0026     for s in segs:
0027         cur.append(s); acc += s[1]
0028         if k < nl - 1 and acc >= cuts[k]:
0029             groups.append(cur); cur = []; k += 1
0030     if cur: groups.append(cur)
0031     for g in groups:
0032         with open(os.path.join(outdir, "list_%02d.txt" % ilist), "w") as f:
0033             for s in g: f.write(s[2] + "\n")
0034         summary.append("%02d  run %d  segments %4d-%4d  nseg %3d  events %8d" % (ilist, r, g[0][0], g[-1][0], len(g), sum(s[1] for s in g)))
0035         ilist += 1
0036 with open(os.path.join(outdir, "lists_summary.txt"), "w") as f:
0037     f.write("\n".join(summary) + "\n")
0038 nev = sum(int(r[2]) for r in rows if int(r[2]) > 0)
0039 with open(os.path.join(outdir, "manifest.txt"), "w") as f:
0040     f.write("collect_dir  %s\ntarget       %d\nnlists       %d\nnfiles       %d\nnevents      %d\nmade         %s\n"
0041             % (cdir, TARGET, ilist, sum(len(v) for v in runs.values()), nev, datetime.datetime.now().isoformat(timespec="seconds")))
0042 print("\n".join(summary))
0043 print("lists: %d   files: %d   events: %d   (unreadable files: %d)" % (ilist, sum(len(v) for v in runs.values()), nev, len(bad)))