Plotting of particle occurences

scripts/nuclei_xsections.py

@@ -22,7 +22,6 @@ import os
 import time
 import numpy as np
 import km3pipe as kp
-import matplotlib.pyplot as plt
 from tempfile import TemporaryDirectory
 from os.path import dirname, abspath
 from collections import defaultdict
@@ -102,23 +101,49 @@ def worker(energy, Z, anti_flag=False):
 
 def plot(datafile, plotfile):
     import h5py
-    f = h5py.File(datafile, 'r')
-    data = f['xsec'][()]
-    data = data[['energy', 'Z', 'sum']]
-    energies = np.sort(np.unique(data['energy']))
-    Zrange = np.sort(np.unique(data['Z']))
-    fig, ax = plt.subplots()
-    h = ax.hist2d(x=data['energy'],
-                  y=data['Z'],
-                  bins=(energies, Zrange),
-                  weights=np.divide(data['sum'], data['energy']))
-    cb = plt.colorbar(h[3], ax=ax)
-    cb.ax.get_yaxis().labelpad = 15
-    cb.ax.set_ylabel(r"$\nu$ Crosssection / E [$10^-38cm^{cm^2}/GeV$]")
-    plt.xlabel("Energy [GeV]")
-    plt.ylabel(r"Z")
-    plt.xscale('log')
-    plt.savefig(plotfile)
+    import matplotlib.pyplot as plt
+    from matplotlib.backends.backend_pdf import PdfPages
+    import matplotlib.colors as colors
+    with PdfPages(plotfile) as pdf:
+        f = h5py.File(datafile, 'r')
+        data = f['xsec'][()]
+        data = data[['energy', 'Z', 'sum']]
+        energies = np.sort(np.unique(data['energy']))
+        Zrange = np.sort(np.unique(data['Z']))
+        fig, ax = plt.subplots()
+        h = ax.hist2d(x=data['energy'],
+                      y=data['Z'],
+                      bins=(energies, Zrange),
+                      weights=np.divide(data['sum'], data['energy']))
+        cb = plt.colorbar(h[3], ax=ax)
+        cb.ax.get_yaxis().labelpad = 15
+        cb.ax.set_ylabel(r"$\nu$ Crosssection / E [$10^-38cm^{cm^2}/GeV$]")
+        plt.xlabel("Energy [GeV]")
+        plt.ylabel(r"Z")
+        plt.xscale('log')
+        pdf.savefig()
+        plt.close()
+        data = f['particles'][()]
+        data = data[['energy', 'Z', 'multiplicity', 'particles']]
+        particle_ids = np.unique(data['particles'])
+        for p in particle_ids:
+            mask = data['particles'] == p
+            fig, ax = plt.subplots()
+            tmp = data[mask]
+            h = ax.hist2d(x=tmp['energy'],
+                          y=tmp['Z'],
+                          bins=(energies, Zrange),
+                          weights=tmp['multiplicity'],
+                          norm=colors.LogNorm())
+            cb = plt.colorbar(h[3], ax=ax)
+            cb.ax.get_yaxis().labelpad = 15
+            cb.ax.set_ylabel(r"Particle Count")
+            plt.xlabel("Energy [GeV]")
+            plt.ylabel(r"Z")
+            plt.xscale('log')
+            plt.title('Particle %s' % p)
+            pdf.savefig()
+            plt.close()
 
 
 class GiBUUTaskPump(kp.Module):
@@ -173,7 +198,7 @@ def main():
     if args['--Zmax']:
         Zmax = int(args['--Zmax'])
 
-    energies = np.logspace(-1, 2, 100)
+    energies = np.logspace(-1, 3, 100)
     Zrange = range(Zmin, Zmax + 1)
     tasks = {}
     with ThreadPoolExecutor(max_workers=workers) as executor: