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