| Geant4 Cross Reference |
1 import argparse 1 import argparse
2 2
3 import numpy as np 3 import numpy as np
4 4
5 from core.constants import INIT_DIR, GEN_DIR, 5 from core.constants import INIT_DIR, GEN_DIR, N_CELLS_PHI, N_CELLS_R, N_CELLS_Z
6 from utils.observables import LongitudinalProf 6 from utils.observables import LongitudinalProfile, LateralProfile, Energy
7 from utils.plotters import ProfilePlotter, Ene 7 from utils.plotters import ProfilePlotter, EnergyPlotter
8 from utils.preprocess import load_showers 8 from utils.preprocess import load_showers
9 9
10 10
11 def parse_args(): 11 def parse_args():
12 p = argparse.ArgumentParser() 12 p = argparse.ArgumentParser()
13 p.add_argument("--geometry", type=str, def 13 p.add_argument("--geometry", type=str, default="")
14 p.add_argument("--energy", type=int, defau 14 p.add_argument("--energy", type=int, default="")
15 p.add_argument("--angle", type=int, defaul 15 p.add_argument("--angle", type=int, default="")
16 args = p.parse_args() 16 args = p.parse_args()
17 return args 17 return args
18 18
19 19
20 # main function 20 # main function
21 def main(): 21 def main():
22 # Parse commandline arguments 22 # Parse commandline arguments
23 args = parse_args() 23 args = parse_args()
24 particle_energy = args.energy 24 particle_energy = args.energy
25 particle_angle = args.angle 25 particle_angle = args.angle
26 geometry = args.geometry 26 geometry = args.geometry
27 # 1. Full simulation data loading 27 # 1. Full simulation data loading
28 # Load energy of showers from a single geo 28 # Load energy of showers from a single geometry, energy and angle
29 e_layer_g4 = load_showers(INIT_DIR, geomet 29 e_layer_g4 = load_showers(INIT_DIR, geometry, particle_energy,
30 particle_angle) 30 particle_angle)
31 # 2. Fast simulation data loading, scaling 31 # 2. Fast simulation data loading, scaling to original energy range & reshaping
32 vae_energies = np.load(f"{GEN_DIR}/VAE_Gen 32 vae_energies = np.load(f"{GEN_DIR}/VAE_Generated_Geo_{geometry}_E_{particle_energy}_Angle_{particle_angle}.npy")
33 # Reshape the events into 3D 33 # Reshape the events into 3D
34 e_layer_vae = vae_energies.reshape((len(va 34 e_layer_vae = vae_energies.reshape((len(vae_energies), N_CELLS_R, N_CELLS_PHI, N_CELLS_Z))
35 35
36 print("Data has been loaded.") 36 print("Data has been loaded.")
37 37
38 # 3. Create observables from raw data. 38 # 3. Create observables from raw data.
39 full_sim_long = LongitudinalProfile(_input 39 full_sim_long = LongitudinalProfile(_input=e_layer_g4)
40 full_sim_lat = LateralProfile(_input=e_lay 40 full_sim_lat = LateralProfile(_input=e_layer_g4)
41 full_sim_energy = Energy(_input=e_layer_g4 41 full_sim_energy = Energy(_input=e_layer_g4)
42 ml_sim_long = LongitudinalProfile(_input=e 42 ml_sim_long = LongitudinalProfile(_input=e_layer_vae)
43 ml_sim_lat = LateralProfile(_input=e_layer 43 ml_sim_lat = LateralProfile(_input=e_layer_vae)
44 ml_sim_energy = Energy(_input=e_layer_vae) 44 ml_sim_energy = Energy(_input=e_layer_vae)
45 45
46 print("Created observables.") 46 print("Created observables.")
47 47
48 # 4. Plot observables 48 # 4. Plot observables
49 longitudinal_profile_plotter = ProfilePlot 49 longitudinal_profile_plotter = ProfilePlotter(particle_energy, particle_angle, geometry, full_sim_long, ml_sim_long,
50 50 _plot_gaussian=False)
51 lateral_profile_plotter = ProfilePlotter(p 51 lateral_profile_plotter = ProfilePlotter(particle_energy, particle_angle,
52 g 52 geometry, full_sim_lat, ml_sim_lat, _plot_gaussian=False)
53 energy_plotter = EnergyPlotter(particle_en 53 energy_plotter = EnergyPlotter(particle_energy, particle_angle, geometry, full_sim_energy, ml_sim_energy)
54 54
55 longitudinal_profile_plotter.plot_and_save 55 longitudinal_profile_plotter.plot_and_save()
56 lateral_profile_plotter.plot_and_save() 56 lateral_profile_plotter.plot_and_save()
57 energy_plotter.plot_and_save() 57 energy_plotter.plot_and_save()
58 print("Done.") 58 print("Done.")
59 59
60 60
61 if __name__ == "__main__": 61 if __name__ == "__main__":
62 exit(main()) 62 exit(main())