Add missing dependencies, add new example

python/examples/particle_histogram.py

+#!/usr/bin/python3
+
+import argparse
+import boost_histogram as bh
+import corsika
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import mplhep
+import numpy as np
+import os
+import pandas as pd
+import particle
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--input-dir", required=True, help="input corsika output library")
+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)
+
+lib = corsika.Library(args.input_dir)
+
+particles_config = lib.get("particles").config
+particles = lib.get("particles").astype("pandas")
+
+hist = bh.Histogram(
+    bh.axis.IntCategory([], growth=True),  # PID
+    bh.axis.Regular(18 * 10, 1e-3, 1e15, transform=bh.axis.transform.log),  # energy
+    bh.axis.Regular(20, 0, 2000),  # lat. distance
+    bh.axis.Regular(360, 0, 360, circular=True),  # azimuth
+)
+
+hist.fill(
+    particles["pdg"],
+    particles["kinetic_energy"],
+    np.hypot(particles["x"], particles["y"]),
+    np.arctan2(particles["y"], particles["x"]),
+    weight=particles["weight"],
+)
+
+fig, ax = plt.subplots(dpi=200)
+ax.set(xscale="log", xlabel="count")
+v = hist[:, bh.sum, bh.sum, bh.sum].values()
+ypos = np.arange(len(v))
+ax.barh(ypos, v)
+ax.set_yticks(ypos)
+ax.set_yticklabels([f"${particle.Particle.from_pdgid(pid).latex_name}$" for pid in hist.axes[0]])
+
+output_name = os.path.join(args.output_dir, "particle_histogram_types.pdf")
+print("Saving", output_name)
+plt.savefig(output_name)
+
+
+fig, ax = plt.subplots(dpi=200)
+ax.set(xscale="log", yscale="log", ylabel="counts", xlabel=f"kinetic energy")
+
+pids = [22, 11, -11, 211, -211, 13, -13, 2212, 2112]
+
+for pid in pids:
+    name = particle.Particle.from_pdgid(pid).latex_name
+    try: # Not all particle types might be in the file
+        mplhep.histplot(hist[bh.loc(pid), :, bh.sum, bh.sum], ax=ax, label=f"${name}$, {pid}")
+    except:
+        pass
+
+ax.legend()
+ax.set_xlim(min(particles["kinetic_energy"]), max(particles["kinetic_energy"]))
+
+output_name = os.path.join(args.output_dir, "particle_histogram_energies.pdf")
+print("Saving", output_name)
+plt.savefig(output_name)

python/setup.py

@@ -46,8 +46,11 @@ setup(
         ],
         "pandas": ["pandas"],
         "examples" : [
+            "argparse",
             "matplotlib",
-            "argparse"
+            "pandas",
+            "particle",
+            "mplhep",
         ]
     },
     scripts=[],