training/utils/preprocess.py

0001 import h5py
0002 import numpy as np
0003
0004 from core.constants import INIT_DIR, ORIGINAL_DIM, MAX_ENERGY, MAX_ANGLE, MIN_ANGLE, MIN_ENERGY
0005
0006
0007 # preprocess function loads the data and returns the array of the shower energies and the condition arrays
0008 def preprocess():
0009     energies_train = []
0010     cond_e_train = []
0011     cond_angle_train = []
0012     cond_geo_train = []
0013     # This example is trained using 2 detector geometries
0014     for geo in ["SiW", "SciPb"]:
0015         dir_geo = INIT_DIR + geo + "/"
0016         # loop over the angles in a step of 10
0017         for angle_particle in range(MIN_ANGLE, MAX_ANGLE + 10, 10):
0018             f_name = f"{geo}_angle_{angle_particle}.h5"
0019             f_name = dir_geo + f_name
0020             # read the HDF5 file
0021             h5 = h5py.File(f_name, "r")
0022             # loop over energies from min_energy to max_energy
0023             energy_particle = MIN_ENERGY
0024             while energy_particle <= MAX_ENERGY:
0025                 # scale the energy of each cell to the energy of the primary particle (in MeV units)
0026                 events = np.array(h5[f"{energy_particle}"]) / (energy_particle * 1000)
0027                 energies_train.append(events.reshape(len(events), ORIGINAL_DIM))
0028                 # build the energy and angle condition vectors
0029                 cond_e_train.append([energy_particle / MAX_ENERGY] * len(events))
0030                 cond_angle_train.append([angle_particle / MAX_ANGLE] * len(events))
0031                 # build the geometry condition vector (1 hot encoding vector)
0032                 if geo == "SiW":
0033                     cond_geo_train.append([[0, 1]] * len(events))
0034                 if geo == "SciPb":
0035                     cond_geo_train.append([[1, 0]] * len(events))
0036                 energy_particle *= 2
0037     # return numpy arrays
0038     energies_train = np.concatenate(energies_train)
0039     cond_e_train = np.concatenate(cond_e_train)
0040     cond_angle_train = np.concatenate(cond_angle_train)
0041     cond_geo_train = np.concatenate(cond_geo_train)
0042     return energies_train, cond_e_train, cond_angle_train, cond_geo_train
0043
0044
0045 # get_condition_arrays function returns condition values from a single geometry, a single energy and angle of primary
0046 # particles
0047 """
0048     - geo : name of the calorimeter geometry (eg: SiW, SciPb)
0049     - energy_particle : energy of the primary particle in GeV units
0050     - nb_events : number of events
0051 """
0052
0053
0054 def get_condition_arrays(geo, energy_particle, nb_events):
0055     cond_e = [energy_particle / MAX_ENERGY] * nb_events
0056     cond_angle = [energy_particle / MAX_ENERGY] * nb_events
0057     if geo == "SiW":
0058         cond_geo = [[0, 1]] * nb_events
0059     else:  # geo == "SciPb"
0060         cond_geo = [[1, 0]] * nb_events
0061     cond_e = np.array(cond_e)
0062     cond_angle = np.array(cond_angle)
0063     cond_geo = np.array(cond_geo)
0064     return cond_e, cond_angle, cond_geo
0065
0066
0067 # load_showers function loads events from a single geometry, a single energy and angle of primary particles
0068 """
0069     - init_dir: the name of the directory which contains the HDF5 files
0070     - geo : name of the calorimeter geometry (eg: SiW, SciPb)
0071     - energy_particle : energy of the primary particle in GeV units
0072     - angle_particle : angle of the primary particle in degrees
0073 """
0074
0075
0076 def load_showers(init_dir, geo, energy_particle, angle_particle):
0077     dir_geo = init_dir + geo + "/"
0078     f_name = f"{geo}_angle_{angle_particle}.h5"
0079     f_name = dir_geo + f_name
0080     # read the HDF5 file
0081     h5 = h5py.File(f_name, "r")
0082     energies = np.array(h5[f"{energy_particle}"])
0083     return energies