Copy over changes from !559

corsika/detail/modules/writers/PrimaryWriter.inl

@@ -41,6 +41,8 @@ namespace corsika {
         static_cast<int>(get_PDG(pID));
     output_["shower_" + std::to_string(showerId_)]["name"] =
         static_cast<std::string>(get_name(pID));
+    output_["shower_" + std::to_string(showerId_)]["total_energy"] =
+        (kineticEnergy + get_mass(pID)) / 1_GeV;
     output_["shower_" + std::to_string(showerId_)]["kinetic_energy"] =
         kineticEnergy / 1_GeV;
     output_["shower_" + std::to_string(showerId_)]["x"] = x / 1_m;

python/Makefile

@@ -16,13 +16,13 @@ tests:
 	${PYTHON} -m pytest --cov=corsika tests
 
 flake:
-	${PYTHON} -m flake8 corsika tests
+	${PYTHON} -m flake8 corsika tests examples
 
 black:
-	${PYTHON} -m black -t py37 corsika tests
+	${PYTHON} -m black -t py37 corsika tests examples
 
 isort:
-	${PYTHON} -m isort --atomic corsika tests
+	${PYTHON} -m isort --atomic corsika tests examples
 
 mypy:
 	${PYTHON} -m mypy corsika

python/README.md

@@ -52,3 +52,9 @@ The example scripts require additional dependencies that can be installed.
 ```shell
 pip install argparse matplotlib particle
 ```
+
+Examples are can be run like this:
+
+```shell
+python examples/shower_profile.py --input-dir <path-to-C8-output>
+```

python/examples/.gitignore

+*.pdf
+*.png
\ No newline at end of file

python/examples/particle_distribution.py

+#!/usr/bin/python3
+
+import argparse
+import os
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+import particle
+
+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(
+    "--output-dir",
+    default=os.path.join(here, "example_plots"),
+    help="output directory for plots",
+)
+args = parser.parse_args()
+
+if not os.path.isdir(args.output_dir):
+    print("Making directory", args.output_dir)
+    os.makedirs(args.output_dir)
+
+# Load the shower simulation output
+lib = corsika.Library(args.input_dir)
+
+# Load the primary particle information from the shower
+primaries = lib.get("primary").data
+primary = primaries[0]
+primary_config = lib.get("primary").config
+
+# Get the contents of the "particles" sub-directory
+particles_config = lib.get("particles").config  # meta information
+particles = lib.get("particles").astype("pandas")  # particle info
+
+# Quick sanity check
+if particles_config["plane"] != primary_config["plane"]:
+    print("Primary", primary_config["plane"])
+    print("Particles on observation plane", particles_config["plane"])
+    msg = "The observation plane of the primary is not the same as that"
+    msg += " of the particle output. This example will not work as intended"
+    msg += " since they do not share the same coordinate system"
+    raise RuntimeError(msg)
+
+r = np.sqrt(particles["x"] ** 2 + particles["y"] ** 2)
+
+if not len(r):
+    print("No particles were found on the observation plane, quitting...")
+    exit()
+
+max_r = max(r)
+
+######################
+# Make a plot of the x-y positions of all particles on the observation plane
+# Note that these are in the coordinate system of the plane
+######################
+
+n_bins = int(np.sqrt(len(r)) * 2)
+bins = np.linspace(-max_r, max_r, n_bins)
+
+length_unit_str = particles_config["units"]["length"]  # read in the printed units
+
+title = f"Primary: {primary.name}," + r" E$_{\rm tot}$:"
+title += f" {primary.total_energy:.2e} {primary_config['units']['energy']},"
+title += f" Zen: {np.rad2deg(np.pi - np.arccos(primary.nz)):0.1f} deg"
+
+fig, ax = plt.subplots(1, 1)
+ax.set_aspect("equal")
+ax.set_title(particles_config["type"] + "\n" + title)
+ax.set_xlabel(f"observation plane x ({length_unit_str})")
+ax.set_ylabel(f"observation plane y ({length_unit_str})")
+
+ax.hist2d(
+    particles["x"],
+    particles["y"],
+    weights=particles["weight"],
+    bins=(bins, bins),
+    norm=mpl.colors.LogNorm(),
+)
+
+plot_path = os.path.join(args.output_dir, "particle_dist.png")
+print("Saving plot", plot_path)
+fig.savefig(plot_path)
+
+######################
+# Make a plot of the lateral distribution of particles
+# this is based off the approximate shower axis
+######################
+
+r_dot_n = (particles["x"] - primary.x) * primary.nx + (
+    particles["y"] - primary.y
+) * primary.ny
+displacement_sq = (particles["x"] - primary.x) ** 2 + (particles["y"] - primary.y) ** 2
+
+if "z" in particles.keys():  # z may not be written out
+    r_dot_n += (particles["z"] - primary.z) * primary.nz
+    displacement_sq += (particles["z"] - primary.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))
+
+fig, ax = plt.subplots(1, 1)
+ax.set_title(title)
+ax.set_xlabel(f"distance to shower axis ({length_unit_str})")
+ax.set_ylabel("number of particles")
+ax.set_yscale("log")
+
+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
+    if pid not in particles["pdg"]:
+        continue
+    plt.hist(
+        r_axis,
+        weights=(particles["pdg"] == pid).astype(float),
+        bins=bins,
+        label=f"${name}$, {pid}",
+        histtype="step",
+    )
+
+ax.legend()
+
+plot_path = os.path.join(args.output_dir, "particle_lateral.png")
+print("Saving plot", plot_path)
+fig.savefig(plot_path)

python/examples/radio_emission.py

+#!/usr/bin/python3
+
+import argparse
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+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(
+    "--output-dir",
+    default=os.path.join(here, "example_plots"),
+    help="output directory for plots",
+)
+args = parser.parse_args()
+
+if not os.path.isdir(args.output_dir):
+    print("Making directory", args.output_dir)
+    os.makedirs(args.output_dir)
+
+# Load the shower simulation output
+lib = corsika.Library(args.input_dir)
+
+# Load the primary particle information from the shower
+primaries = lib.get("primary").data
+n_showers = len(primaries)
+primary_config = lib.get("primary").config
+
+found_zhs = False
+try:
+    waveforms_config_zhs = lib.get("ZHS").config  # meta information
+    waveforms_zhs = lib.get("ZHS").astype("pandas")
+    antennas_zhs = lib.get("ZHS").get_antennas()
+    found_zhs = True
+except Exception:
+    print("No ZHS waveforms in this directory")
+
+
+found_coreas = False
+try:
+    waveforms_config_coreas = lib.get("CoREAS").config  # meta information
+    waveforms_coreas = lib.get("CoREAS").astype("pandas")
+    antennas_coreas = lib.get("CoREAS").get_antennas()
+    found_coreas = True
+except Exception:
+    print("No CoREAS waveforms in this directory")
+
+if not found_zhs and not found_coreas:
+    print("No electric field waveforms are in this file, quitting...")
+    exit()
+
+for ishower in range(n_showers):
+
+    primary = primaries[ishower]
+
+    title = f"Primary: {primary.name}," + r" E$_{\rm tot}$:"
+    title += f" {primary.total_energy:.2e} {primary_config['units']['energy']}"
+
+    # Get all of the unique antenna names to match up the CoREAS/ZHS waveforms
+    ant_names = []
+    if found_zhs:
+        ant_names += list(waveforms_zhs[str(ishower)].keys())
+        zhs_dict = waveforms_zhs[str(ishower)]
+    if found_coreas:
+        ant_names += list(waveforms_coreas[str(ishower)].keys())
+        coreas_dict = waveforms_coreas[str(ishower)]
+    ant_names = np.unique(ant_names)
+
+    ncols = 1
+    nrows = len(ant_names)
+    fig, ax = plt.subplots(
+        ncols=ncols,
+        nrows=nrows,
+        figsize=(ncols * 8, nrows * 3),
+        gridspec_kw={"wspace": 0.2, "hspace": 0.3},
+    )
+
+    # Make a plot for each of the antenna locations
+    for iant, ant_name in enumerate(ant_names):
+
+        # Add blank entry to show colors in legend
+        ax[iant].fill_between([], [], [], color="k", label="Ex")
+        ax[iant].fill_between([], [], [], color="r", label="Ey")
+        ax[iant].fill_between([], [], [], color="b", label="Ez")
+
+        if not iant:
+            ax[iant].set_title(title)
+
+        # Make the plots for the ZHS waveforms
+        if found_zhs and ant_name in antennas_zhs.keys():
+            ant_position = np.array(
+                [
+                    antennas_zhs[ant_name]["x"],
+                    antennas_zhs[ant_name]["y"],
+                    antennas_zhs[ant_name]["z"],
+                ]
+            ).flatten()
+            times = np.array(zhs_dict[ant_name]["time"])
+
+            ax[iant].plot(
+                times, zhs_dict[ant_name]["Ex"], color="k", linestyle="-", label="ZHS"
+            )
+            ax[iant].plot(times, zhs_dict[ant_name]["Ey"], color="r", linestyle="-")
+            ax[iant].plot(times, zhs_dict[ant_name]["Ez"], color="b", linestyle="-")
+            ax[iant].set_xlabel(f"Time [{waveforms_config_zhs['units']['time']}]")
+            ax[iant].set_ylabel(
+                f"Amp ({waveforms_config_zhs['units']['electric field']})"
+            )
+
+        # Make the plots for the CoREAS waveforms
+        if found_coreas and ant_name in antennas_coreas.keys():
+            ant_position = np.array(
+                [
+                    antennas_coreas[ant_name]["x"],
+                    antennas_coreas[ant_name]["y"],
+                    antennas_coreas[ant_name]["z"],
+                ]
+            ).flatten()
+            times = np.array(coreas_dict[ant_name]["time"])
+
+            ax[iant].plot(
+                times,
+                coreas_dict[ant_name]["Ex"],
+                color="k",
+                linestyle="--",
+                label="CoREAS",
+            )
+            ax[iant].plot(times, coreas_dict[ant_name]["Ey"], color="r", linestyle="--")
+            ax[iant].plot(times, coreas_dict[ant_name]["Ez"], color="b", linestyle="--")
+            ax[iant].set_xlabel(f"Time ({waveforms_config_coreas['units']['time']})")
+            ax[iant].set_ylabel(
+                f"Amp ({waveforms_config_coreas['units']['electric field']})"
+            )
+
+        ymin, ymax = ax[iant].get_ylim()
+        max_amp = max(abs(ymin), abs(ymax))
+        ax[iant].set_ylim(-max_amp, max_amp)
+        text = f"Antenna @ ({ant_position[0]:0.0f}, {ant_position[1]:0.0f},"
+        text += f" {ant_position[2]:0.0f})"
+        ax[iant].text(times[-1], 2 * max_amp * 0.05 - max_amp, text, ha="right")
+        ax[iant].legend(loc="upper right")
+
+    plot_path = os.path.join(args.output_dir, f"AntennaWaveforms_Sh{ishower}.png")
+    print("Saving", plot_path)
+    fig.savefig(plot_path, bbox_inches="tight")

python/examples/shower_profile.py

+#!/usr/bin/python3
+
+import argparse
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+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(
+    "--output-dir",
+    default=os.path.join(here, "example_plots"),
+    help="output directory for plots",
+)
+args = parser.parse_args()
+
+if not os.path.isdir(args.output_dir):
+    print("Making directory", args.output_dir)
+    os.makedirs(args.output_dir)
+
+
+def plot_avg_profile(dat, part, ax):
+    # Plots the average profile for particle type `part`
+    # averages over all showers in the simulation output dir
+    nshower = len(dat.shower.unique())  # number of showers in output dir
+    max_dX = np.max(dat["X"])
+    h = np.histogram(
+        dat.X, bins=np.linspace(0, max_dX, 200), weights=dat[part] * 1 / float(nshower)
+    )
+    ax.plot(h[1][:-1], h[0], label=part)
+
+
+# Load the shower simulation output
+lib = corsika.Library(args.input_dir)
+
+# Load the primary particle information from the shower
+primaries = lib.get("primary").data
+primary = primaries[0]
+primary_config = lib.get("primary").config
+
+# Get the contents of the "profile" sub-directory
+pr_config = lib.get("profile").config  # meta information
+pr = lib.get("profile").astype("pandas")  # particle-number values
+
+# Get the contents of the "energyloss" sub-directory
+prdx_config = lib.get("energyloss").config  # meta information
+prdx = lib.get("energyloss").astype("pandas")  # energy loss spectrum
+
+title = f"Primary: {primary.name}," + r" E$_{\rm tot}$:"
+title += f" {primary.total_energy:.2e} {primary_config['units']['energy']},"
+title += f" Zen: {np.rad2deg(np.pi - np.arccos(primary.nz)):0.1f} deg"
+
+
+def draw_profiles(pr, pr_config):
+    # Plots the number of particles as a function of slant depth
+    # individual distributions are made for each particle type
+    fig, ax = plt.subplots(1, 1)
+    for part in pr.keys():
+        # Skip shower id number and slant depth columns
+        if "shower" == part or "X" == part:
+            continue
+        plot_avg_profile(pr, part, ax)
+    ax.set_title(pr_config["type"] + "\n" + title)
+    unit_grammage_str = pr_config["units"]["grammage"]  # get units of simulation output
+    ax.set_xlabel(f"slant depth, X ({unit_grammage_str})")
+    ax.set_ylabel("dN/dX")
+    ax.legend()
+    ax.set_yscale("log")
+    ax.set_xlim(min(pr["X"]), max(pr["X"]))
+
+    plot_path = os.path.join(args.output_dir, "show_profile_profiles.png")
+    print("Saving", plot_path)
+    fig.savefig(plot_path)
+
+
+def draw_energyloss(prdx, prdx_config):
+    # Plots the energy deposition as a function of slant depth
+    fig, ax = plt.subplots(1, 1)
+    ax.set_title(prdx_config["type"] + "\n" + title)
+    plot_avg_profile(prdx, "total", ax)
+    unit_energy_str = prdx_config["units"]["energy"]  # get units of simulation output
+    unit_grammage_str = prdx_config["units"]["grammage"]
+    ax.set_xlabel(f"slant depth, X ({unit_grammage_str})")
+    ax.set_ylabel(f"dE/dX ({unit_energy_str} / ({unit_grammage_str}))")
+    ax.set_xlim(min(prdx["X"]), max(prdx["X"]))
+
+    plot_path = os.path.join(args.output_dir, "show_profile_energyloss.png")
+    print("Saving", plot_path)
+    fig.savefig(plot_path)
+
+
+draw_profiles(pr, pr_config)
+draw_energyloss(prdx, prdx_config)

python/setup.py

@@ -45,6 +45,11 @@ setup(
             "types-PyYAML",
             "pandas-stubs",
         ],
+        "examples": [
+            "argparse",
+            "matplotlib",
+            "particle",
+        ],
     },
     scripts=[],
     project_urls={"code": "https://gitlab.iap.kit.edu/AirShowerPhysics/corsika"},