add function to print particle probabilities in plot-muons

2 files changed, 29 insertions, 0 deletions

diff --git a/utils/plot-muons b/utils/plot-muons
index d069d01..0646e2e 100755
--- a/utils/plot-muons
+++ b/utils/plot-muons
@@ -28,6 +28,20 @@ from scipy.stats import iqr, poisson
 from scipy.stats import iqr, norm, beta
 from sddm.stats import *
 
+particle_id = {20: 'e', 22: 'u'}
+
+def print_particle_probs(data):
+    n = [len(data[data.id == id]) for id in (20,22,2020,2022,2222)]
+
+    alpha = np.ones_like(n) + n
+
+    mode = dirichlet_mode(alpha)
+    std = np.sqrt(dirichlet.var(alpha))
+
+    for i, id in enumerate((20,22,2020,2022,2222)):
+        particle_id_str = ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)])
+        print("P(%s) = %.2f +/- %.2f" % (particle_id_str,mode[i]*100,std[i]*100))
+
 if __name__ == '__main__':
     import argparse
     import numpy as np
@@ -111,6 +125,12 @@ if __name__ == '__main__':
     # cause the reconstruction to overestimate the energy.
     michel_low_nhit = michel[michel.muon_nhit < 2500]
 
+    print("Particle ID probability for Michel electrons:")
+    print("Data")
+    print_particle_probs(michel_low_nhit)
+    print("Monte Carlo")
+    print_particle_probs(michel_mc)
+
     fig = plt.figure()
     plot_hist2_data_mc(michel_low_nhit,michel_mc)
     despine(fig,trim=True)
@@ -175,6 +195,12 @@ if __name__ == '__main__':
 
         print(stopping_muons[['run','gtid','ke','T_dx','dT','gtid_michel','r_michel','ftp_r_michel','id','r']])
 
+        print("Particle ID probability for Stopping Muons:")
+        print("Data")
+        print_particle_probs(stopping_muons)
+        print("Monte Carlo")
+        print_particle_probs(stopping_muons_mc)
+
         fig = plt.figure()
         plot_hist2_data_mc(stopping_muons,stopping_muons_mc)
         despine(fig,trim=True)
diff --git a/utils/sddm/stats.py b/utils/sddm/stats.py
index 9c5e827..01bb009 100644
--- a/utils/sddm/stats.py
+++ b/utils/sddm/stats.py
@@ -41,6 +41,9 @@ def get_multinomial_prob(data, mc, norm, size=10000):
     chi2_data = np.array(chi2_data)
     return np.count_nonzero(chi2_samples >= chi2_data)/len(chi2_samples)
 
+def dirichlet_mode(alpha):
+    return (alpha-1)/(np.sum(alpha)-len(alpha))
+
 def plot_hist2_data_mc(data, mc, bins=None, norm=None):
     import matplotlib.pyplot as plt
     from matplotlib.lines import Line2D