diff --git a/utils/__pycache__/bib_pdgs.cpython-310.pyc b/utils/__pycache__/bib_pdgs.cpython-310.pyc new file mode 100644 index 0000000..e03d042 Binary files /dev/null and b/utils/__pycache__/bib_pdgs.cpython-310.pyc differ diff --git a/utils/fluka_to_edm4hep.py b/utils/fluka_to_edm4hep.py new file mode 100755 index 0000000..f20e8c1 --- /dev/null +++ b/utils/fluka_to_edm4hep.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python +"""This script converts a FLUKA binary file to an SLCIO file with LCIO::MCParticle instances""" + +import os +import argparse +import numpy as np + + +parser = argparse.ArgumentParser(description='Convert FLUKA binary file to SLCIO file with MCParticles') +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('-c', '--comment', metavar='TEXT', help='Comment to be added to the header', type=str) +parser.add_argument('-b', '--bx_time', metavar='TIME', help='Time of the bunch crossing [s]', type=float, default=0.0) +parser.add_argument('-n', '--normalization', metavar='N', help='Normalization of the generated sample', type=float, default=1.0) +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='+') +parser.add_argument('--nopdgs', metavar='ID', help='PDG IDs of particles to be excluded', type=int, default=None, nargs='+') +parser.add_argument('--ne_min', metavar='E', help='Minimum energy of accepted neutrons [GeV]', type=float, default=None) +parser.add_argument('--t_max', metavar='T', help='Maximum time of accepted particles [ns]', type=float, default=None) + +args = parser.parse_args() + +if not args.overwrite and os.path.isfile(args.file_out): + raise FileExistsError(f'Output file already exists: {args.file_out:s}') + + +from math import sqrt +from pdb import set_trace as br +from array import array + +from edm4hep import edm4hep +from ROOT import podio +from podio.root_io import Writer +import cppyy + + +import random +import math + +from bib_pdgs import FLUKA_PIDS, PDG_PROPS + +def bytes_from_file(filename): + with open(filename, 'rb') as f: + while True: + chunk = np.fromfile(f, dtype=line_dt, count=1) + if not len(chunk): + return + yield chunk + +# Binary format of a single entry +line_dt=np.dtype([ + ('fid', np.int32), + ('fid_mo', np.int32), + ('E', np.float64), + ('x', np.float64), + ('y', np.float64), + ('z', np.float64), + ('cx', np.float64), + ('cy', np.float64), + ('cz', np.float64), + ('age', np.float64), + ('age_mu', np.float64), + ('x_mu', np.float64), + ('y_mu', np.float64), + ('z_mu', np.float64), + ('x_mo', np.float64), + ('y_mo', np.float64), + ('z_mo', np.float64), + ('px_mo', np.float64), + ('py_mo', np.float64), + ('pz_mo', np.float64), + ('age_mo', np.float64) +]) + +######################################## Start of the processing +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}') + +# Initialize the EDM4HEP file writer +writer = Writer(args.file_out) + +# Write a RunHeader +frame = podio.Frame() +frame.putParameter("InputFiles", len(args.files_in)) +frame.putParameter("Normalization", str(args.normalization)) +frame.putParameter("BXTime", str(args.bx_time)) +frame.putParameter("FilesPerEvent", str(args.files_event)) + +if args.t_max: + frame.putParameter("Time_max", str(args.t_max)) +if args.ne_min: + frame.putParameter("NeutronEnergy_min", str(args.ne_min)) +if args.pdgs: + frame.putParameter("PdgIds", str(args.pdgs)) +if args.nopdgs: + frame.putParameter("NoPdgIds", str(args.nopdgs)) +if args.comment: + frame.putParameter("Comment", str(args.comment)) + +writer.writeFrame(frame, 'header') + +# Bookkeeping variables +random.seed() +nEventFiles = 0 +nLines = 0 +nEvents = 0 +col = None +evt = None + +# Reading the complete files +for iF, file_in in enumerate(args.files_in): + # Creating the EDM4HEP event and collection + if nEventFiles == 0: + col = edm4hep.MCParticleCollection() + evt = podio.Frame() + evt.putParameter("eventNumber", str(nEvents)) + + + # 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) * 1e9 + + # 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 = edm4hep.MutableMCParticle() + particle = col.create() + 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)): + # 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 + particle.setVertex(pos) + particle.setMomentum(mom) + # Adding particle to the collection + # col.push_back(particle) + + # Updating counters + nEventFiles += 1 + if nEventFiles >= args.files_event or iF+1 == len(args.files_in): + nEvents +=1 + nEventFiles = 0 + evt.put(cppyy.gbl.std.move(col), "MCParticles") + writer.writeFrame(evt, 'events') + print(f'Wrote event: {nEvents:d} with {col.size()} particles') + +print(f'Wrote {nEvents:d} events to file: {args.file_out:s}') + +writer.finish() \ No newline at end of file