diff --git a/utils/fluka_to_slcio.py b/utils/fluka_to_slcio.py index d3f3703..ea9f9db 100644 --- a/utils/fluka_to_slcio.py +++ b/utils/fluka_to_slcio.py @@ -7,11 +7,11 @@ parser = argparse.ArgumentParser(description='Convert FLUKA binary file to SLCIO file with MCParticles') -parser.add_argument('file_in', metavar='FILE_IN', help='Path to the input FLUKA file') -parser.add_argument('file_out', metavar='FILE_OUT.slcio', help='Path to the output SLCIO file') +parser.add_argument('files_in', metavar='FILE_IN', help='Input binary FLUKA file(s)', nargs='+') +parser.add_argument('file_out', metavar='FILE_OUT.slcio', help='Output SLCIO file') parser.add_argument('-b', '--bx_time', metavar='TIME', help='Time of the bunch crossing [ns]', type=float, default=117.78143152396451) parser.add_argument('-n', '--normalization', metavar='N', help='Normalization of the generated sample', type=float, default=1.0) -parser.add_argument('-l', '--lines_event', metavar='L', help='Number of lines to put in a single LCIO event (default: 1000)', type=int, default=1000) +parser.add_argument('-f', '--files_event', metavar='L', help='Number of files to merge into a single LCIO event (default: 1)', type=int, default=1) parser.add_argument('-m', '--max_lines', metavar='M', help='Maximum number of lines to process', type=int, default=None) parser.add_argument('-o', '--overwrite', help='Overwrite existing output file', action='store_true', default=False) parser.add_argument('--pdgs', metavar='ID', help='PDG IDs of particles to be included', type=int, default=None, nargs='+') @@ -69,8 +69,8 @@ def bytes_from_file(filename): ]) ######################################## Start of the processing -print(f'Converting data from file\n {args.file_in:s}\nto SLCIO file\n {args.file_out:s}\nwith normalization: {args.normalization:.1f}') -print(f'Splitting by {args.lines_event:d} lines/event'); +print(f'Converting data from {len(args.files_in)} file(s)\nto SLCIO file: {args.file_out:s}\nwith normalization: {args.normalization:.1f}') +print(f'Storing {args.files_event:d} files/event'); if args.pdgs is not None: print(f'Will only use particles with PDG IDs: {args.pdgs}') @@ -81,7 +81,7 @@ def bytes_from_file(filename): # Write a RunHeader run = IMPL.LCRunHeaderImpl() run.setRunNumber(0) -run.parameters().setValue('InputPath', args.file_in) +run.parameters().setValue('InputPath', str(args.files_in)) run.parameters().setValue('Normalization', args.normalization) if args.t_max: run.parameters().setValue('Time_max', args.t_max) @@ -95,106 +95,103 @@ def bytes_from_file(filename): # Bookkeeping variables random.seed() -nEventLines = 0 +nEventFiles = 0 +nLines = 0 nEvents = 0 col = None evt = None -# Reading the the file one entry at a time -for iL, data in enumerate(bytes_from_file(args.file_in)): - if nEventLines == 0: +# Reading the complete files +for iF, file_in in enumerate(args.files_in): + # Creating the LCIO event and collection + if nEventFiles == 0: col = IMPL.LCCollectionVec(EVENT.LCIO.MCPARTICLE) - # Creating the LCIO event and collection evt = IMPL.LCEventImpl() evt.setEventNumber(nEvents) evt.addCollection(col, 'MCParticle') - if args.max_lines and iL >= args.max_lines: - break - - # Extracting relevant values from the line - fid,e, x,y,z, cx,cy,cz, toff,toff_mo = (data[n][0] for n in [ - 'fid', 'E', - 'x','y','z', - 'cx', 'cy', 'cz', - 'age', 'age_mo' - ]) - - # Converting FLUKA ID to PDG ID - try: - pdg = FLUKA_PIDS[fid] - except KeyError: - print(f'WARNING: Unknown PDG ID for FLUKA ID: {pid}') - continue - - # Calculating the absolute time of the particle [ns] - t = toff - toff_mo - args.bx_time - - # Skipping if particle's time is greater than allowed - if args.t_max is not None and t > args.t_max: - continue - - # Calculating the components of the momentum vector - mom = np.array([cx, cy, cz], dtype=np.float32) - mom *= e - - # Skipping if it's a neutron with too low kinetic energy - if args.ne_min is not None and abs(pdg) == 2112 and np.linalg.norm(mom) < args.ne_min: - continue - - # Getting the charge and mass of the particle - if pdg not in PDG_PROPS: - print('WARNING! No properties defined for PDG ID: {0:d}'.format(pdg)) - print(' Skpping the particle...') - continue - charge, mass = PDG_PROPS[pdg] - - # Calculating how many random copies of the particle to create according to the weight - nP_frac, nP = math.modf(args.normalization) - if nP_frac > 0 and random.random() < nP_frac: - nP += 1 - nP = int(nP) - - # Creating the particle with original parameters - particle = IMPL.MCParticleImpl() - particle.setPDG(pdg) - particle.setGeneratorStatus(1) - particle.setTime(t) - particle.setMass(mass) - particle.setCharge(charge) - pos = np.array([x, y, z], dtype=np.float64) - - # Creating the particle copies with random Phi rotation - px, py, pz = mom - for i, iP in enumerate(range(nP)): - p = IMPL.MCParticleImpl(particle) - # Rotating position and momentum of the copies by a random angle in Phi - if i > 0: - dPhi = random.random() * math.pi * 2 - co = math.cos(dPhi) - si = math.sin(dPhi) - pos[0] = co * x - si * y - pos[1] = si * x + co * y - mom[0] = co * px - si * py - mom[1] = si * px + co * py - p.setVertex(pos) - p.setMomentum(mom) - # Adding particle to the collection - col.addElement(p) + # Looping over particles from the file + for iL, data in enumerate(bytes_from_file(file_in)): + if args.max_lines and nLines >= args.max_lines: + break + + # Extracting relevant values from the line + fid,e, x,y,z, cx,cy,cz, toff,toff_mo = (data[n][0] for n in [ + 'fid', 'E', + 'x','y','z', + 'cx', 'cy', 'cz', + 'age', 'age_mo' + ]) + + # Converting FLUKA ID to PDG ID + try: + pdg = FLUKA_PIDS[fid] + except KeyError: + print(f'WARNING: Unknown PDG ID for FLUKA ID: {fid}') + continue + + # Calculating the absolute time of the particle [ns] + t = toff - toff_mo - args.bx_time + + # Skipping if particle's time is greater than allowed + if args.t_max is not None and t > args.t_max: + continue + + # Calculating the components of the momentum vector + mom = np.array([cx, cy, cz], dtype=np.float32) + mom *= e + + # Skipping if it's a neutron with too low kinetic energy + if args.ne_min is not None and abs(pdg) == 2112 and np.linalg.norm(mom) < args.ne_min: + continue + + # Getting the charge and mass of the particle + if pdg not in PDG_PROPS: + print('WARNING! No properties defined for PDG ID: {0:d}'.format(pdg)) + print(' Skpping the particle...') + continue + charge, mass = PDG_PROPS[pdg] + + # Calculating how many random copies of the particle to create according to the weight + nP_frac, nP = math.modf(args.normalization) + if nP_frac > 0 and random.random() < nP_frac: + nP += 1 + nP = int(nP) + + # Creating the particle with original parameters + particle = IMPL.MCParticleImpl() + particle.setPDG(pdg) + particle.setGeneratorStatus(1) + particle.setTime(t) + particle.setMass(mass) + particle.setCharge(charge) + pos = np.array([x, y, z], dtype=np.float64) + + # Creating the particle copies with random Phi rotation + px, py, pz = mom + for i, iP in enumerate(range(nP)): + p = IMPL.MCParticleImpl(particle) + # Rotating position and momentum of the copies by a random angle in Phi + if i > 0: + dPhi = random.random() * math.pi * 2 + co = math.cos(dPhi) + si = math.sin(dPhi) + pos[0] = co * x - si * y + pos[1] = si * x + co * y + mom[0] = co * px - si * py + mom[1] = si * px + co * py + p.setVertex(pos) + p.setMomentum(mom) + # Adding particle to the collection + col.addElement(p) # Updating counters - nEventLines += 1 - if nEventLines >= args.lines_event: + nEventFiles += 1 + if nEventFiles >= args.files_event or iF+1 == len(args.files_in): nEvents +=1 - nEventLines = 0 + nEventFiles = 0 wrt.writeEvent(evt) - print('Wrote event: {0:d} (line {1:d})'.format(nEvents, iL+1)) + print(f'Wrote event: {nEvents:d} with {col.getNumberOfElements()} particles') -# Writing the last event -if nEventLines > 0: - wrt.writeEvent(evt) - nEvents += 1 - print('Wrote event: {0:d}'.format(nEvents)) - -print('Wrote {0:d} events to file: {1:s}'.format(nEvents, args.file_out)) +print(f'Wrote {nEvents:d} events to file: {args.file_out:s}') wrt.close()