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