| Geant4 Cross Reference |
1 import h5py 1 import h5py
2 import numpy as np 2 import numpy as np
3 3
4 from core.constants import INIT_DIR, ORIGINAL_ 4 from core.constants import INIT_DIR, ORIGINAL_DIM, MAX_ENERGY, MAX_ANGLE, MIN_ANGLE, MIN_ENERGY
5 5
6 6
7 # preprocess function loads the data and retur 7 # preprocess function loads the data and returns the array of the shower energies and the condition arrays
8 def preprocess(): 8 def preprocess():
9 energies_train = [] 9 energies_train = []
10 cond_e_train = [] 10 cond_e_train = []
11 cond_angle_train = [] 11 cond_angle_train = []
12 cond_geo_train = [] 12 cond_geo_train = []
13 # This example is trained using 2 detector 13 # This example is trained using 2 detector geometries
14 for geo in ["SiW", "SciPb"]: 14 for geo in ["SiW", "SciPb"]:
15 dir_geo = INIT_DIR + geo + "/" 15 dir_geo = INIT_DIR + geo + "/"
16 # loop over the angles in a step of 10 16 # loop over the angles in a step of 10
17 for angle_particle in range(MIN_ANGLE, 17 for angle_particle in range(MIN_ANGLE, MAX_ANGLE + 10, 10):
18 f_name = f"{geo}_angle_{angle_part 18 f_name = f"{geo}_angle_{angle_particle}.h5"
19 f_name = dir_geo + f_name 19 f_name = dir_geo + f_name
20 # read the HDF5 file 20 # read the HDF5 file
21 h5 = h5py.File(f_name, "r") 21 h5 = h5py.File(f_name, "r")
22 # loop over energies from min_ener 22 # loop over energies from min_energy to max_energy
23 energy_particle = MIN_ENERGY 23 energy_particle = MIN_ENERGY
24 while energy_particle <= MAX_ENERG 24 while energy_particle <= MAX_ENERGY:
25 # scale the energy of each cel 25 # scale the energy of each cell to the energy of the primary particle (in MeV units)
26 events = np.array(h5[f"{energy 26 events = np.array(h5[f"{energy_particle}"]) / (energy_particle * 1000)
27 energies_train.append(events.r 27 energies_train.append(events.reshape(len(events), ORIGINAL_DIM))
28 # build the energy and angle c 28 # build the energy and angle condition vectors
29 cond_e_train.append([energy_pa 29 cond_e_train.append([energy_particle / MAX_ENERGY] * len(events))
30 cond_angle_train.append([angle 30 cond_angle_train.append([angle_particle / MAX_ANGLE] * len(events))
31 # build the geometry condition 31 # build the geometry condition vector (1 hot encoding vector)
32 if geo == "SiW": 32 if geo == "SiW":
33 cond_geo_train.append([[0, 33 cond_geo_train.append([[0, 1]] * len(events))
34 if geo == "SciPb": 34 if geo == "SciPb":
35 cond_geo_train.append([[1, 35 cond_geo_train.append([[1, 0]] * len(events))
36 energy_particle *= 2 36 energy_particle *= 2
37 # return numpy arrays 37 # return numpy arrays
38 energies_train = np.concatenate(energies_t 38 energies_train = np.concatenate(energies_train)
39 cond_e_train = np.concatenate(cond_e_train 39 cond_e_train = np.concatenate(cond_e_train)
40 cond_angle_train = np.concatenate(cond_ang 40 cond_angle_train = np.concatenate(cond_angle_train)
41 cond_geo_train = np.concatenate(cond_geo_t 41 cond_geo_train = np.concatenate(cond_geo_train)
42 return energies_train, cond_e_train, cond_ 42 return energies_train, cond_e_train, cond_angle_train, cond_geo_train
43 43
44 44
45 # get_condition_arrays function returns condit 45 # get_condition_arrays function returns condition values from a single geometry, a single energy and angle of primary
46 # particles 46 # particles
47 """ 47 """
48 - geo : name of the calorimeter geometry ( 48 - geo : name of the calorimeter geometry (eg: SiW, SciPb)
49 - energy_particle : energy of the primary 49 - energy_particle : energy of the primary particle in GeV units
50 - nb_events : number of events 50 - nb_events : number of events
51 """ 51 """
52 52
53 53
54 def get_condition_arrays(geo, energy_particle, 54 def get_condition_arrays(geo, energy_particle, nb_events):
55 cond_e = [energy_particle / MAX_ENERGY] * 55 cond_e = [energy_particle / MAX_ENERGY] * nb_events
56 cond_angle = [energy_particle / MAX_ENERGY 56 cond_angle = [energy_particle / MAX_ENERGY] * nb_events
57 if geo == "SiW": 57 if geo == "SiW":
58 cond_geo = [[0, 1]] * nb_events 58 cond_geo = [[0, 1]] * nb_events
59 else: # geo == "SciPb" 59 else: # geo == "SciPb"
60 cond_geo = [[1, 0]] * nb_events 60 cond_geo = [[1, 0]] * nb_events
61 cond_e = np.array(cond_e) 61 cond_e = np.array(cond_e)
62 cond_angle = np.array(cond_angle) 62 cond_angle = np.array(cond_angle)
63 cond_geo = np.array(cond_geo) 63 cond_geo = np.array(cond_geo)
64 return cond_e, cond_angle, cond_geo 64 return cond_e, cond_angle, cond_geo
65 65
66 66
67 # load_showers function loads events from a si 67 # load_showers function loads events from a single geometry, a single energy and angle of primary particles
68 """ 68 """
69 - init_dir: the name of the directory whic 69 - init_dir: the name of the directory which contains the HDF5 files
70 - geo : name of the calorimeter geometry ( 70 - geo : name of the calorimeter geometry (eg: SiW, SciPb)
71 - energy_particle : energy of the primary 71 - energy_particle : energy of the primary particle in GeV units
72 - angle_particle : angle of the primary pa 72 - angle_particle : angle of the primary particle in degrees
73 """ 73 """
74 74
75 75
76 def load_showers(init_dir, geo, energy_particl 76 def load_showers(init_dir, geo, energy_particle, angle_particle):
77 dir_geo = init_dir + geo + "/" 77 dir_geo = init_dir + geo + "/"
78 f_name = f"{geo}_angle_{angle_particle}.h5 78 f_name = f"{geo}_angle_{angle_particle}.h5"
79 f_name = dir_geo + f_name 79 f_name = dir_geo + f_name
80 # read the HDF5 file 80 # read the HDF5 file
81 h5 = h5py.File(f_name, "r") 81 h5 = h5py.File(f_name, "r")
82 energies = np.array(h5[f"{energy_particle} 82 energies = np.array(h5[f"{energy_particle}"])
83 return energies 83 return energies