diff --git a/corsika/detail/modules/writers/PrimaryWriter.inl b/corsika/detail/modules/writers/PrimaryWriter.inl index 600f40ec3082feee1a6843844cf91493b6d36306..84046db7d56adcde3c19f68322deea9efdc4a5be 100644 --- a/corsika/detail/modules/writers/PrimaryWriter.inl +++ b/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; diff --git a/python/Makefile b/python/Makefile index 54716dd66ec1955d351445c123d90200fae23ffb..787077a1945bc91016039ead46e24f674957894f 100644 --- a/python/Makefile +++ b/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 diff --git a/python/README.md b/python/README.md index 12789324a2d1ba073664a7aeaef83098f6091dc4..ede02b1274111c6246c0a4dcf59ab0ecfe337d0d 100644 --- a/python/README.md +++ b/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> +``` diff --git a/python/examples/.gitignore b/python/examples/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..4d23a7a189e36c95d9907b652302ed23c7bf927f --- /dev/null +++ b/python/examples/.gitignore @@ -0,0 +1,2 @@ +*.pdf +*.png \ No newline at end of file diff --git a/python/examples/particle_distribution.py b/python/examples/particle_distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..87b5772bb8d2d12ff09ce95ed51ed23cbafd9f34 --- /dev/null +++ b/python/examples/particle_distribution.py @@ -0,0 +1,133 @@ +#!/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) diff --git a/python/examples/radio_emission.py b/python/examples/radio_emission.py new file mode 100644 index 0000000000000000000000000000000000000000..f65791cee60a4d69dfdb64870ed82db7974c494f --- /dev/null +++ b/python/examples/radio_emission.py @@ -0,0 +1,152 @@ +#!/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") diff --git a/python/examples/shower_profile.py b/python/examples/shower_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..3f4b1e6a70e7f0b83a6526b282385c967fa2c73b --- /dev/null +++ b/python/examples/shower_profile.py @@ -0,0 +1,100 @@ +#!/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) diff --git a/python/setup.py b/python/setup.py index 20da6eaa249c572540bd66ba574a4995705d5578..72005141bd998347b865fef1d45dcaf09bcc3f14 100644 --- a/python/setup.py +++ b/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"},