Gabriele Cosmo
110 lines
3.4 KiB
Python
110 lines
3.4 KiB
Python
"""
|
|
Converts a ROOT file to an HDF5 file, saving the shower energy in a 3D array.
|
|
|
|
Args:
|
|
file_name: Path to the ROOT file with a name:
|
|
output_<NAME>_angle<ANGLE>_<NUM>events_fullSim_<ID>.root
|
|
output_dir: Path to the output directory, will create a file
|
|
output_dir/<NAME>_Angle_<ANGLE>_<NUM>showers_<ID>.h5
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
#!/bin/env python
|
|
import sys
|
|
import argparse
|
|
import numpy as np
|
|
import os
|
|
import uproot
|
|
import h5py
|
|
|
|
# Number of cells in the r, phi & z directions
|
|
NB_CELLS_R = 18
|
|
NB_CELLS_PHI = 50
|
|
NB_CELLS_Z = 45
|
|
|
|
|
|
def find_between(s, first, last):
|
|
"""
|
|
Find a substring between two other substrings.
|
|
|
|
Args:
|
|
s (str): The string to search in.
|
|
first (str): The first substring.
|
|
last (str): The last substring.
|
|
|
|
Returns:
|
|
str: The substring found between 'first' and 'last'.
|
|
"""
|
|
try:
|
|
start = s.index(first) + len(first)
|
|
end = s.index(last, start)
|
|
return s[start:end]
|
|
except ValueError:
|
|
return ""
|
|
|
|
|
|
def parse_args(argv):
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--outDir", type=str, default="")
|
|
p.add_argument("--fName", type=str, default="")
|
|
args = p.parse_args()
|
|
return args
|
|
|
|
|
|
def main(argv):
|
|
# Parse commandline arguments
|
|
args = parse_args(argv)
|
|
file_name = args.fName
|
|
output_dir = args.outDir
|
|
# The energy and angle of the particle are part
|
|
# of the ROOT's file name so they can be extracted from the name
|
|
# Get the energy value of the particle
|
|
energy_particle = find_between(file_name, "output_", "_angle")
|
|
# Get the angle value of the particule
|
|
angle_particle = find_between(file_name, "_angle", "_")
|
|
# Get the number of showers
|
|
num_showers = find_between(file_name, "_", "events_")
|
|
# Get Ids specific to the generated files
|
|
file_id = find_between(file_name, "fullSim_", ".root")
|
|
if os.stat(file_name).st_size > 0:
|
|
h5_file = h5py.File(
|
|
f"{output_dir}/{energy_particle}_Angle_{angle_particle}_{num_showers}showers_{file_id}.h5", "w"
|
|
)
|
|
# Read the Root file
|
|
file = uproot.open(file_name)
|
|
energy_particle = file["global"]["EnergyMC"].array()
|
|
cell_r = file["virtualReadout"]["rhoCell"].array()
|
|
cell_phi = file["virtualReadout"]["phiCell"].array()
|
|
cell_energy = file["virtualReadout"]["EnergyCell"].array()
|
|
cell_z = file["virtualReadout"]["zCell"].array()
|
|
all_events = []
|
|
# loop over events
|
|
for event in range(len(energy_particle)):
|
|
# Initialize a 3D array with shape nb_events, nb_cells in x,y,z
|
|
data = np.zeros((NB_CELLS_R, NB_CELLS_PHI, NB_CELLS_Z))
|
|
for ind in range(len(cell_r[event])):
|
|
# This if statement is added to avoid having extra un-indexed cells
|
|
if (
|
|
(cell_r[event][ind] < NB_CELLS_R)
|
|
and (cell_phi[event][ind] < NB_CELLS_PHI)
|
|
and (cell_z[event][ind] < NB_CELLS_Z)
|
|
):
|
|
data[cell_r[event][ind]][cell_phi[event][ind]][
|
|
cell_z[event][ind]
|
|
] = cell_energy[event][ind]
|
|
all_events.append(data)
|
|
# Save dataset
|
|
h5_file.create_dataset(
|
|
f"{energy_particle}",
|
|
data=all_events,
|
|
compression="gzip",
|
|
compression_opts=9,
|
|
)
|
|
h5_file.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main(sys.argv[1:]))
|