Gabriele Cosmo
58 lines
2.7 KiB
Python
58 lines
2.7 KiB
Python
"""
|
|
** utils **
|
|
defines the data loading and preprocessing function
|
|
"""
|
|
|
|
# Setup
|
|
import h5py
|
|
import numpy as np
|
|
|
|
# preprocess function returns the array of the shower energies and the condition arrays
|
|
"""
|
|
- init_dir: the name of the directory which contains the HDF5 files
|
|
- size_1DVec: represents the size of the input and output layer of the VAE which corresponds to the total number of readout cells
|
|
- min_energy,max_energy: minimum and maximum primary particle energy to consider for training in GeV units
|
|
- min_angle and max_angle: minimum and maximum primary particle angle to consider for training in degrees units
|
|
"""
|
|
def preprocess(init_dir,size_1DVec,min_angle,max_angle,min_energy,max_energy):
|
|
energies_Train = []
|
|
condE_Train = []
|
|
condAngle_Train = []
|
|
condGeo_Train = []
|
|
# This example is trained using 2 detector geometries
|
|
for geo in [ 'SiW' , 'SciPb' ]:
|
|
dirGeo = init_dir + geo + '/'
|
|
energyParticle=min_energy
|
|
# loop over the energies in powers of 2
|
|
while(energyParticle<=max_energy):
|
|
# loop over the angles in a step of 10
|
|
for angleParticle in range(min_angle,max_angle+10,10):
|
|
fName = 'Energy_%s_Angle_%s.hdf5' %(energyParticle,angleParticle)
|
|
fName = dirGeo + fName
|
|
# read the HDF5 file
|
|
h5 = h5py.File(fName,'r')
|
|
# get the key value of the group from the HDF5 file
|
|
GroupKey = 'Grp_Angle_%s_E_%s'%(angleParticle,energyParticle)
|
|
# get all key values of one group
|
|
listKeys = list( h5[GroupKey].keys() )
|
|
# loop over the events
|
|
for ckey in listKeys:
|
|
# scale the energy of each cell to the energy of the primary particle (in MeV units)
|
|
energyArray = np.array(h5[GroupKey][ckey])/(energyParticle*1000)
|
|
energies_Train.append( energyArray.reshape(size_1DVec) )
|
|
# build the energy and angle condition vectors
|
|
condE_Train.append( [energyParticle/mamax_energyxE]*len(listKeys) )
|
|
condAngle_Train.append( [angleParticle/max_angle]*len(listKeys) )
|
|
# build the geometry condition vector (1 hot encoding vector)
|
|
if( geo == 'SiW' ):
|
|
condGeo_Train.append( [[0,1]]*len(listKeys) )
|
|
else:
|
|
condGeo_Train.append( [[1,0]]*len(listKeys) )
|
|
energyParticle*=2
|
|
# return numpy arrays
|
|
energies_Train = np.array(energies_Train)
|
|
condE_Train = np.concatenate(condE_Train)
|
|
condAngle_Train = np.concatenate(condAngle_Train)
|
|
condGeo_Train = np.concatenate(condGeo_Train)
|
|
return energies_Train,condE_Train,condAngle_Train,condGeo_Train
|