diff --git a/python/examples/first_interactions.py b/python/examples/first_interactions.py index ef2635638b9bee35ec036fdf3a487855649dffa6..8df6c37e593fd25827dfc71893b3b69c984c5c06 100644 --- a/python/examples/first_interactions.py +++ b/python/examples/first_interactions.py @@ -12,9 +12,7 @@ import corsika here = os.path.abspath(os.path.dirname(__file__)) parser = argparse.ArgumentParser() -parser.add_argument( - "--input-dir", required=True, help="output directory of the CORSIKA 8 simulation" -) +parser.add_argument("--input-dir", required=True, help="output directory of the CORSIKA 8 simulation") parser.add_argument( "--output-dir", default=os.path.join(here, "example_plots"), @@ -49,7 +47,7 @@ for sh_id in shower_ids: data = interactions[interactions["shower"] == sh_id] mother = data[data["primary"]] - norm = np.sqrt(mother["px"]**2 + mother["py"]**2 + mother["pz"]**2) + norm = np.sqrt(mother["px"] ** 2 + mother["py"] ** 2 + mother["pz"] ** 2) angle = np.arctan2(mother["py"].iloc[0], mother["px"].iloc[0]) pz = mother["pz"].iloc[0] px = np.cos(-angle) * mother["px"].iloc[0] + np.sin(-angle) * mother["py"].iloc[0] @@ -58,14 +56,11 @@ for sh_id in shower_ids: name = f"${particle.Particle.from_pdgid(pdg).latex_name}$" ax[sh_id].plot([-px, 0], [-pz, 0], label=name) - # Get daughters, filter out zero-momentum particles daughters = data[data["primary"] == False] - p_tot = np.sqrt(daughters["px"]**2 + daughters["py"]**2 + daughters["pz"]**2) + p_tot = np.sqrt(daughters["px"] ** 2 + daughters["py"] ** 2 + daughters["pz"] ** 2) daughters = daughters[p_tot > 0] - print(daughters) - for i in range(len(daughters)): pz = daughters["pz"].iloc[i] px = np.cos(-angle) * daughters["px"].iloc[i] + np.sin(-angle) * daughters["py"].iloc[i] @@ -75,8 +70,7 @@ for sh_id in shower_ids: name = f"${particle.Particle.from_pdgid(pdg).latex_name}$" except Exception: name = "Unknown" - print(pdg) - + ax[sh_id].plot([px, 0], [pz, 0], label=name) ax[sh_id].set_aspect("equal")