Gabriele Cosmo
63 lines
2.4 KiB
Python
63 lines
2.4 KiB
Python
import argparse
|
|
|
|
import numpy as np
|
|
|
|
from core.constants import INIT_DIR, GEN_DIR, N_CELLS_PHI, N_CELLS_R, N_CELLS_Z
|
|
from utils.observables import LongitudinalProfile, LateralProfile, Energy
|
|
from utils.plotters import ProfilePlotter, EnergyPlotter
|
|
from utils.preprocess import load_showers
|
|
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--geometry", type=str, default="")
|
|
p.add_argument("--energy", type=int, default="")
|
|
p.add_argument("--angle", type=int, default="")
|
|
args = p.parse_args()
|
|
return args
|
|
|
|
|
|
# main function
|
|
def main():
|
|
# Parse commandline arguments
|
|
args = parse_args()
|
|
particle_energy = args.energy
|
|
particle_angle = args.angle
|
|
geometry = args.geometry
|
|
# 1. Full simulation data loading
|
|
# Load energy of showers from a single geometry, energy and angle
|
|
e_layer_g4 = load_showers(INIT_DIR, geometry, particle_energy,
|
|
particle_angle)
|
|
# 2. Fast simulation data loading, scaling to original energy range & reshaping
|
|
vae_energies = np.load(f"{GEN_DIR}/VAE_Generated_Geo_{geometry}_E_{particle_energy}_Angle_{particle_angle}.npy")
|
|
# Reshape the events into 3D
|
|
e_layer_vae = vae_energies.reshape((len(vae_energies), N_CELLS_R, N_CELLS_PHI, N_CELLS_Z))
|
|
|
|
print("Data has been loaded.")
|
|
|
|
# 3. Create observables from raw data.
|
|
full_sim_long = LongitudinalProfile(_input=e_layer_g4)
|
|
full_sim_lat = LateralProfile(_input=e_layer_g4)
|
|
full_sim_energy = Energy(_input=e_layer_g4)
|
|
ml_sim_long = LongitudinalProfile(_input=e_layer_vae)
|
|
ml_sim_lat = LateralProfile(_input=e_layer_vae)
|
|
ml_sim_energy = Energy(_input=e_layer_vae)
|
|
|
|
print("Created observables.")
|
|
|
|
# 4. Plot observables
|
|
longitudinal_profile_plotter = ProfilePlotter(particle_energy, particle_angle, geometry, full_sim_long, ml_sim_long,
|
|
_plot_gaussian=False)
|
|
lateral_profile_plotter = ProfilePlotter(particle_energy, particle_angle,
|
|
geometry, full_sim_lat, ml_sim_lat, _plot_gaussian=False)
|
|
energy_plotter = EnergyPlotter(particle_energy, particle_angle, geometry, full_sim_energy, ml_sim_energy)
|
|
|
|
longitudinal_profile_plotter.plot_and_save()
|
|
lateral_profile_plotter.plot_and_save()
|
|
energy_plotter.plot_and_save()
|
|
print("Done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|