File indexing completed on 2026-08-31 08:24:38
0001
0002 """
0003 make_mbdtrkz_lists.py <runnumber> [--nevents N] [--trkseg S] [--trkdst clus|tracks]
0004 [--listdir DIR]
0005
0006 Writes clus.list or tracks.list (controlled by --trkdst), and fit.list covering
0007 N events starting at block S (each block is N events).
0008
0009 --nevents N events per block (default 1000, must be a multiple of
0010 EVT_PER_SEED=1000)
0011 --trkdst clus|tracks source DST type: 'clus' reads from dst_clus/<runnumber>.list
0012 and writes clus.list; 'tracks' reads from dst_tracks/<runnumber>.list
0013 and writes tracks.list (default: tracks)
0014 --listdir DIR base directory containing dst_clus/, dst_tracks/, dst_fit/
0015 (default: ~/sphenix/sphenix_bbc/run2025/lists/run3pp)
0016
0017 Example:
0018 --nevents 100000 --trkdst clus -> clus segs 0-99, fit seg 0
0019 --nevents 100000 --trkdst tracks -> tracks segs 100-199, fit seg 1
0020 --nevents 200000 -> tracks segs 200-399, fit segs 2-3
0021
0022 Returns 0 on success, 1 if any expected DSTs are missing.
0023 """
0024
0025 import os
0026 import re
0027 import sys
0028
0029 EVT_PER_SEED = 1000
0030 EVT_PER_FIT = 100000
0031 SEEDS_PER_FIT = EVT_PER_FIT // EVT_PER_SEED
0032
0033
0034 def get_seg(filename):
0035 """Return the 5-digit trailing segment number from a DST filename, or None."""
0036 m = re.search(r'-(\d{5})\.root', filename)
0037 return int(m.group(1)) if m else None
0038
0039
0040 def lines_in_range(listfile, lo, hi):
0041 """Return stripped non-empty lines whose segment number is in [lo, hi]."""
0042 result = []
0043 with open(listfile) as fh:
0044 for line in fh:
0045 line = line.strip()
0046 if not line:
0047 continue
0048 seg = get_seg(line)
0049 if seg is not None and lo <= seg <= hi:
0050 result.append(line)
0051 return result
0052
0053
0054 def parse_args():
0055 args = sys.argv[1:]
0056 runnumber = None
0057 nevents = EVT_PER_FIT
0058 trkseg = 0
0059 trkdst = 'tracks'
0060 listdir = os.path.expanduser('~/sphenix/sphenix_bbc/run2025/lists/run3pp')
0061 i = 0
0062 while i < len(args):
0063 if args[i] in ('--nevents', '-nevents'):
0064 if i + 1 >= len(args):
0065 print('ERROR: --nevents requires an argument', file=sys.stderr)
0066 sys.exit(1)
0067 nevents = int(args[i + 1])
0068 i += 2
0069 elif args[i] in ('--trkseg', '-trkseg'):
0070 if i + 1 >= len(args):
0071 print('ERROR: --trkseg requires an argument', file=sys.stderr)
0072 sys.exit(1)
0073 trkseg = int(args[i + 1])
0074 i += 2
0075 elif args[i] in ('--trkdst', '-trkdst'):
0076 if i + 1 >= len(args):
0077 print('ERROR: --trkdst requires an argument', file=sys.stderr)
0078 sys.exit(1)
0079 trkdst = args[i + 1]
0080 if trkdst not in ('clus', 'tracks'):
0081 print(f'ERROR: --trkdst must be "clus" or "tracks", got "{trkdst}"',
0082 file=sys.stderr)
0083 sys.exit(1)
0084 i += 2
0085 elif args[i] in ('--listdir', '-listdir'):
0086 if i + 1 >= len(args):
0087 print('ERROR: --listdir requires an argument', file=sys.stderr)
0088 sys.exit(1)
0089 listdir = os.path.expanduser(args[i + 1])
0090 i += 2
0091 elif not args[i].startswith('-'):
0092 if runnumber is not None:
0093 print('ERROR: unexpected argument', file=sys.stderr)
0094 sys.exit(1)
0095 runnumber = int(args[i])
0096 i += 1
0097 else:
0098 print(f'ERROR: unknown option {args[i]}', file=sys.stderr)
0099 sys.exit(1)
0100
0101 if runnumber is None:
0102 print(f'Usage: {sys.argv[0]} <runnumber> [--nevents N] [--trkseg S] [--trkdst clus|tracks] [--listdir DIR]',
0103 file=sys.stderr)
0104 sys.exit(1)
0105
0106 if nevents <= 0 or nevents % EVT_PER_SEED != 0:
0107 print(f'ERROR: --nevents must be a positive multiple of {EVT_PER_SEED}',
0108 file=sys.stderr)
0109 sys.exit(1)
0110
0111 return runnumber, nevents, trkseg, trkdst, listdir
0112
0113
0114 def main():
0115 runnumber, nevents, trkseg, trkdst, listdir = parse_args()
0116
0117 tracks_per_block = nevents // EVT_PER_SEED
0118 fits_per_block = (nevents + EVT_PER_FIT - 1) // EVT_PER_FIT
0119
0120 src_dir = 'dst_clus' if trkdst == 'clus' else 'dst_tracks'
0121 out_name = 'clus.list' if trkdst == 'clus' else 'tracks.list'
0122
0123 trk_list = os.path.join(listdir, src_dir, f'{runnumber}.list')
0124 fit_list = os.path.join(listdir, 'dst_fit', f'{runnumber}.list')
0125
0126 for f in (trk_list, fit_list):
0127 if not os.path.isfile(f):
0128 print(f'ERROR: list file not found: {f}', file=sys.stderr)
0129 sys.exit(1)
0130
0131
0132 trk_start = trkseg
0133 trk_end = trk_start + tracks_per_block - 1
0134 fit_lo = trk_start // SEEDS_PER_FIT
0135 fit_hi = trk_end // SEEDS_PER_FIT
0136
0137 trk_lines = lines_in_range(trk_list, trk_start, trk_end)
0138 fit_lines = lines_in_range(fit_list, fit_lo, fit_hi)
0139
0140 trk_segs = len({get_seg(l) for l in trk_lines})
0141 fit_segs = len({get_seg(l) for l in fit_lines})
0142
0143 rc = 0
0144 if trk_segs != tracks_per_block:
0145 print(f'ERROR: need {tracks_per_block} {trkdst} segments [{trk_start},{trk_end}], '
0146 f'found {trk_segs}', file=sys.stderr)
0147 rc = 1
0148 if fit_segs != fits_per_block:
0149 print(f'ERROR: need {fits_per_block} fit file(s) for segments [{fit_lo},{fit_hi}], '
0150 f'found {fit_segs}', file=sys.stderr)
0151 rc = 1
0152 if rc:
0153 sys.exit(1)
0154
0155 with open(out_name, 'w') as f:
0156 f.write('\n'.join(trk_lines) + '\n')
0157 with open('fit.list', 'w') as f:
0158 f.write('\n'.join(fit_lines) + '\n')
0159
0160 print(f'{out_name}: {len(trk_lines)} files ({trk_segs} unique segments)')
0161 print(f'fit.list: {len(fit_lines)} files (segs {fit_lo}-{fit_hi})')
0162
0163
0164 if __name__ == '__main__':
0165 main()
