Par04/training/root2h5.py

0001 """
0002 Converts a ROOT file to an HDF5 file, saving the shower energy in a 3D array.
0003
0004 Args:
0005     file_name: Path to the ROOT file with a name:
0006                 output_<NAME>_angle<ANGLE>_<NUM>events_fullSim_<ID>.root
0007     output_dir: Path to the output directory, will create a file
0008                 output_dir/<NAME>_Angle_<ANGLE>_<NUM>showers_<ID>.h5
0009
0010 Returns:
0011     None
0012 """
0013
0014 #!/bin/env python
0015 import sys
0016 import argparse
0017 import numpy as np
0018 import os
0019 import uproot
0020 import h5py
0021
0022 # Number of cells in the r, phi & z directions
0023 NB_CELLS_R = 18
0024 NB_CELLS_PHI = 50
0025 NB_CELLS_Z = 45
0026
0027
0028 def find_between(s, first, last):
0029     """
0030     Find a substring between two other substrings.
0031
0032     Args:
0033         s (str): The string to search in.
0034         first (str): The first substring.
0035         last (str): The last substring.
0036
0037     Returns:
0038         str: The substring found between 'first' and 'last'.
0039     """
0040     try:
0041         start = s.index(first) + len(first)
0042         end = s.index(last, start)
0043         return s[start:end]
0044     except ValueError:
0045         return ""
0046
0047
0048 def parse_args(argv):
0049     p = argparse.ArgumentParser()
0050     p.add_argument("--outDir", type=str, default="")
0051     p.add_argument("--fName", type=str, default="")
0052     args = p.parse_args()
0053     return args
0054
0055
0056 def main(argv):
0057     # Parse commandline arguments
0058     args = parse_args(argv)
0059     file_name = args.fName
0060     output_dir = args.outDir
0061     # The energy and angle of the particle are part
0062     # of the ROOT's file name so they can be extracted from the name
0063     # Get the energy value of the particle
0064     energy_particle = find_between(file_name, "output_", "_angle")
0065     # Get the angle value of the particule
0066     angle_particle = find_between(file_name, "_angle", "_")
0067     # Get the number of showers
0068     num_showers = find_between(file_name, "_", "events_")
0069     # Get Ids specific to the generated files
0070     file_id = find_between(file_name, "fullSim_", ".root")
0071     if os.stat(file_name).st_size > 0:
0072         h5_file = h5py.File(
0073             f"{output_dir}/{energy_particle}_Angle_{angle_particle}_{num_showers}showers_{file_id}.h5", "w"
0074         )
0075         # Read the Root file
0076         file = uproot.open(file_name)
0077         energy_particle = file["global"]["EnergyMC"].array()
0078         cell_r = file["virtualReadout"]["rhoCell"].array()
0079         cell_phi = file["virtualReadout"]["phiCell"].array()
0080         cell_energy = file["virtualReadout"]["EnergyCell"].array()
0081         cell_z = file["virtualReadout"]["zCell"].array()
0082         all_events = []
0083         # loop over events
0084         for event in range(len(energy_particle)):
0085             # Initialize a 3D array with shape nb_events, nb_cells in x,y,z
0086             data = np.zeros((NB_CELLS_R, NB_CELLS_PHI, NB_CELLS_Z))
0087             for ind in range(len(cell_r[event])):
0088                 # This if statement is added to avoid having extra un-indexed cells
0089                 if (
0090                     (cell_r[event][ind] < NB_CELLS_R)
0091                     and (cell_phi[event][ind] < NB_CELLS_PHI)
0092                     and (cell_z[event][ind] < NB_CELLS_Z)
0093                 ):
0094                     data[cell_r[event][ind]][cell_phi[event][ind]][
0095                         cell_z[event][ind]
0096                     ] = cell_energy[event][ind]
0097             all_events.append(data)
0098         # Save dataset
0099         h5_file.create_dataset(
0100             f"{energy_particle}",
0101             data=all_events,
0102             compression="gzip",
0103             compression_opts=9,
0104         )
0105     h5_file.close()
0106
0107
0108 if __name__ == "__main__":
0109     exit(main(sys.argv[1:]))