Add additional plot

python/examples/particle_dist.py

@@ -7,6 +7,7 @@ import argparse
 import numpy as np
 import corsika
 import os
+import particle
 
 here = os.path.abspath(os.path.dirname(__file__))
 
@@ -40,7 +41,7 @@ bins = np.linspace(-max_r, max_r, n_bins)
 unit_str = particles_config["units"]["length"]  # read in the printed units
 
 plt.figure()
-plt.axes().set_aspect('equal')
+plt.axes().set_aspect("equal")
 plt.title(particles_config["type"])
 plt.xlabel(f"observation plane x ({unit_str})")
 plt.ylabel(f"observation plane y ({unit_str})")
@@ -48,3 +49,44 @@ plt.ylabel(f"observation plane y ({unit_str})")
 plt.hist2d(particles["x"], particles["y"], weights=particles["weight"], bins=(bins, bins), norm=mpl.colors.LogNorm())
 
 plt.savefig(os.path.join(args.output_dir, "particle_dist.png"))
+
+######################
+## Make a plot of the lateral distribution of particles
+## this is based off the approximate shower axis
+######################
+
+# This is only the approximate trajectory of the shower
+avg_x_dir = np.average(particles["nx"])
+avg_y_dir = np.average(particles["ny"])
+avg_z_dir = np.average(particles["nz"])
+avg_x = np.average(particles["x"])
+avg_y = np.average(particles["y"])
+
+r_dot_n = (particles["x"] - avg_x) * avg_x_dir + (particles["y"] - avg_y) * avg_y_dir
+displacement_sq = (particles["x"] - avg_x) ** 2 + (particles["y"] - avg_y) ** 2
+
+if "z" in particles.keys():  # z may not be written out
+    avg_z = np.average(particles["z"])
+    r_dot_n += (particles["z"] - avg_z) * avg_z_dir
+    displacement_sq += (particles["z"] - avg_z) ** 2
+
+r_axis = np.sqrt(displacement_sq - r_dot_n**2)
+
+bins = np.linspace(0, max(r_axis) * 1.05, int(n_bins / 2 + 1))
+
+plt.figure()
+plt.xlabel(f"distance to shower axis ({unit_str})")
+plt.ylabel("number of particles")
+plt.semilogy()
+
+pids = [22, 11, -11, 211, -211, 13, -13, 2212, 2112]  # List of common EAS particles
+
+for pid in pids:
+    name = particle.Particle.from_pdgid(pid).latex_name
+    try:  # Not all particle types might be in the file
+        plt.hist(r_axis, weights=(particles["pdg"] == pid).astype(float), bins=bins, label=f"${name}$, {pid}", histtype="step")
+    except:
+        pass
+
+plt.legend()
+plt.savefig(os.path.join(args.output_dir, "particle_lateral.png"))