| Geant4 Cross Reference |
1 """
2 Converts a ROOT file to an HDF5 file, saving the shower energy in a 3D array.
3
4 Args:
5 file_name: Path to the ROOT file with a name:
6 output_<NAME>_angle<ANGLE>_<NUM>events_fullSim_<ID>.root
7 output_dir: Path to the output directory, will create a file
8 output_dir/<NAME>_Angle_<ANGLE>_<NUM>showers_<ID>.h5
9
10 Returns:
11 None
12 """
13
14 #!/bin/env python
15 import sys
16 import argparse
17 import numpy as np
18 import os
19 import uproot
20 import h5py
21
22 # Number of cells in the r, phi & z directions
23 NB_CELLS_R = 18
24 NB_CELLS_PHI = 50
25 NB_CELLS_Z = 45
26
27
28 def find_between(s, first, last):
29 """
30 Find a substring between two other substrings.
31
32 Args:
33 s (str): The string to search in.
34 first (str): The first substring.
35 last (str): The last substring.
36
37 Returns:
38 str: The substring found between 'first' and 'last'.
39 """
40 try:
41 start = s.index(first) + len(first)
42 end = s.index(last, start)
43 return s[start:end]
44 except ValueError:
45 return ""
46
47
48 def parse_args(argv):
49 p = argparse.ArgumentParser()
50 p.add_argument("--outDir", type=str, default="")
51 p.add_argument("--fName", type=str, default="")
52 args = p.parse_args()
53 return args
54
55
56 def main(argv):
57 # Parse commandline arguments
58 args = parse_args(argv)
59 file_name = args.fName
60 output_dir = args.outDir
61 # The energy and angle of the particle are part
62 # of the ROOT's file name so they can be extracted from the name
63 # Get the energy value of the particle
64 energy_particle = find_between(file_name, "output_", "_angle")
65 # Get the angle value of the particule
66 angle_particle = find_between(file_name, "_angle", "_")
67 # Get the number of showers
68 num_showers = find_between(file_name, "_", "events_")
69 # Get Ids specific to the generated files
70 file_id = find_between(file_name, "fullSim_", ".root")
71 if os.stat(file_name).st_size > 0:
72 h5_file = h5py.File(
73 f"{output_dir}/{energy_particle}_Angle_{angle_particle}_{num_showers}showers_{file_id}.h5", "w"
74 )
75 # Read the Root file
76 file = uproot.open(file_name)
77 energy_particle = file["global"]["EnergyMC"].array()
78 cell_r = file["virtualReadout"]["rhoCell"].array()
79 cell_phi = file["virtualReadout"]["phiCell"].array()
80 cell_energy = file["virtualReadout"]["EnergyCell"].array()
81 cell_z = file["virtualReadout"]["zCell"].array()
82 all_events = []
83 # loop over events
84 for event in range(len(energy_particle)):
85 # Initialize a 3D array with shape nb_events, nb_cells in x,y,z
86 data = np.zeros((NB_CELLS_R, NB_CELLS_PHI, NB_CELLS_Z))
87 for ind in range(len(cell_r[event])):
88 # This if statement is added to avoid having extra un-indexed cells
89 if (
90 (cell_r[event][ind] < NB_CELLS_R)
91 and (cell_phi[event][ind] < NB_CELLS_PHI)
92 and (cell_z[event][ind] < NB_CELLS_Z)
93 ):
94 data[cell_r[event][ind]][cell_phi[event][ind]][
95 cell_z[event][ind]
96 ] = cell_energy[event][ind]
97 all_events.append(data)
98 # Save dataset
99 h5_file.create_dataset(
100 f"{energy_particle}",
101 data=all_events,
102 compression="gzip",
103 compression_opts=9,
104 )
105 h5_file.close()
106
107
108 if __name__ == "__main__":
109 exit(main(sys.argv[1:]))