create file edm4hep

generation/bib/fluka_to_edm4hep.py

@@ -0,0 +1,207 @@
+#!/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)
+
+	# 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()

generation/pgun/pgun_edm4hep.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+"""
+This script generates a file with LCIO::MCParticles of specified parameters
+
+Most --arguments support three kinds of inputs:
+ * `value` - the same value will be used for each particle
+ * `min max` - a random value in the range [min; max) will be used for each particle
+ * `0 mean sigma` - a random value sampled from a gaussian distribution will be used for each particle
+"""
+
+import os
+import argparse
+
+parser = argparse.ArgumentParser(description='Generate LCIO::MCParticles with specified parameters')
+parser.add_argument('output', metavar='FILE_OUT.slcio', help='Output LCIO file')
+parser.add_argument('-c', '--comment', metavar='TEXT',  help='Comment to be added to the run header', type=str)
+parser.add_argument('-e', '--events', metavar='N', type=int, default=1,  help='Generate N events')
+parser.add_argument('-p', '--particles', metavar='N', type=int, default=1,  help='Generate N particles/event')
+parser.add_argument('-o', '--overwrite', action='store_true',  help='Overwrite existing output file')
+parser.add_argument('--pdg', metavar='ID', type=int, default=[13], nargs='+',  help='PdgIds of the allowed particles')
+parser.add_argument('--dt', metavar='V', type=float, nargs='*', default=0,  help='Time offset [ns]')
+parser.add_argument('--dz', metavar='V', type=float, nargs='*', default=0,  help='Vertex position along Z [mm]')
+parser.add_argument('--d0', metavar='V', type=float, nargs='*', default=0,  help='Vertex position along R [mm]')
+parser.add_argument('--pt', metavar='V', type=float, nargs='*',  help='Tranverse momentum [GeV]')
+parser.add_argument('--p', metavar='V', type=float, nargs='*',  help='Total momentum [GeV]')
+parser.add_argument('--theta', metavar='A', type=float, default=90, nargs='+',  help='Polar angle [deg]')
+
+args = parser.parse_args()
+
+from edm4hep import edm4hep
+from ROOT import podio
+from podio.root_io import Writer
+import cppyy
+
+from pdgs import PDG_PROPS
+from array import array
+import numpy as np
+import math
+
+# Validating the arguments
+if not args.overwrite and os.path.isfile(args.output):
+	raise FileExistsError(f'Output file already exists: {args.output:s}')
+if (args.pt is None and args.p is None) or (args.pt is not None and args.p is not None):
+	raise RuntimeError('Exactly one of --pt or --p has to be specified')
+for pdg in args.pdg:
+	if pdg not in PDG_PROPS:
+		raise RuntimeError(f'Particle properties not defined for pdgId: {pdg}')
+
+
+# Generating sampling distributions for each property (1 value/event)
+sample_size = args.events
+rng = np.random.default_rng(12345)
+samples = {}
+configs = {
+	'dt': args.dt,
+	'dz': args.dz,
+	'd0': args.d0,
+	'theta': args.theta
+}
+if args.pt is not None:
+	configs['pt'] = args.pt
+else:
+	configs['p'] = args.p
+for name, values in configs.items():
+	if values is None:
+		continue
+	if not isinstance(values, list):
+		samples[name] = np.ones(sample_size) * values
+	elif len(values) == 1:
+		samples[name] = np.ones(sample_size) * values[0]
+	elif len(values) == 2:
+		samples[name] = rng.random(sample_size) * (values[1] - values[0]) + values[0]
+	elif len(values) == 3:
+		samples[name] = np.random.normal(values[1], values[2], sample_size)
+# Adding randomised phi angle for d0
+samples['dphi'] = rng.random(sample_size) * math.pi * 2.
+
+# Opening the output file
+writer = Writer(args.output)
+print(f'Opening output file: {args.output}')
+
+# Writing the run headers
+frame = podio.Frame()
+frame.putParameter('pdgIds', str(args.pdg))
+frame.putParameter('events', str(args.events))
+frame.putParameter('particles/event', str(args.particles))
+if args.comment:
+	frame.putParameter('comment', args.comment)
+for name, values in configs.items():
+	header = str(values) if isinstance(values, list) else values
+	frame.putParameter(name, str(header))
+# wrt.writeRunHeader(run)
+writer.writeFrame(frame, 'header')
+
+# Setting counters
+n_events = 0
+n_particles = 0
+
+# Choosing pdgId of each particle randomly if # of pdgIds is different from # of particles/event
+n_pdgs = len(args.pdg)
+choose_random_pdg = True if args.particles != n_pdgs else False
+# Creating actual particles
+for e in range(args.events):
+	col = edm4hep.MCParticleCollection()
+	evt = podio.Frame()
+	evt.putParameter("eventNumber", str(e))
+
+	for p in range(args.particles):
+		pdg_idx = p
+		if choose_random_pdg:
+			pdg_idx = np.random.choice(n_pdgs, 1)[0]
+		pdg = args.pdg[pdg_idx]
+		# Calculating all properties for this particle in the event
+		phi = rng.random() * math.pi * 2.
+		theta = samples['theta'][e]
+		# Calculating momentum vector
+		if 'pt' in configs:
+			pt = samples['pt'][e]
+			px = pt * math.cos(phi)
+			py = pt * math.sin(phi)
+			pz = pt / math.tan(theta)
+		elif 'p' in configs:
+			p = samples['p'][e]
+			px = p * math.cos(phi) * math.sin(theta)
+			py = p * math.sin(phi) * math.sin(theta)
+			pz = p * math.cos(theta)
+		momentum = array('f', [px, py, pz])
+		# Calculating vertex position
+		vx = samples['d0'][e] / 10.0 * math.cos(samples['dphi'][e])
+		vy = samples['d0'][e] / 10.0 * math.sin(samples['dphi'][e])
+		vz = samples['dz'][e] / 10.0
+		vtx = array('d', [vx, vy, vz])
+		# Assigning properties to the MCParticle
+		mcp = col.create()
+		mcp.setGeneratorStatus(1)
+		mcp.setMass(PDG_PROPS[pdg][1])
+		mcp.setCharge(PDG_PROPS[pdg][0])
+		mcp.setPDG(pdg)
+		mcp.setMomentum(momentum)
+		mcp.setVertex(vtx)
+		# Adding particle to the event
+		n_particles += 1
+	# Writing the event
+	n_events += 1
+	if n_events % (args.events / 10) == 0:
+		print(f'Wrote event {n_events}/{args.events}')
+	evt.put(cppyy.gbl.std.move(col), "MCParticles")	
+	writer.writeFrame(evt, 'events')
+# Closing the output file
+writer.finish()
+print(f'Wrote {n_particles} partiles in {n_events} events to file: {args.output}')
+