Import Geant4 11.4.0.beta source tree

2025-06-26 09:17:29 +02:00
parent 20a218bbe1
commit a499fb82e9
1941 changed files with 203285 additions and 95593 deletions
@@ -0,0 +1,62 @@
+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())