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