update plot-fit-results to make nicer plots

1 files changed, 114 insertions, 28 deletions

diff --git a/utils/plot-fit-results b/utils/plot-fit-results
index f95333f..8bdb7ed 100755
--- a/utils/plot-fit-results
+++ b/utils/plot-fit-results
@@ -16,24 +16,10 @@
 
 from __future__ import print_function, division
 import numpy as np
-from scipy.stats import iqr, norm
+from scipy.stats import iqr, norm, beta
 from matplotlib.lines import Line2D
 import itertools
 
-# on retina screens, the default plots are way too small
-# by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
-# Qt5 will scale everything using the dpi in ~/.Xresources
-import matplotlib
-matplotlib.use("Qt5Agg")
-
-matplotlib.rc('font', size=22)
-
-font = {'family':'serif', 'serif': ['computer modern roman']}
-matplotlib.rc('font',**font)
-
-matplotlib.rc('text', usetex=True)
-
-
 IDP_E_MINUS  =    20
 IDP_MU_MINUS =    22
 
@@ -97,6 +83,42 @@ def iqr_std(x):
         return np.nan
     return iqr(x.values)/1.3489795
 
+def quantile_error(x,q):
+    """
+    Returns the standard error for the qth quantile of `x`. The error is
+    computed using the Maritz-Jarrett method described here:
+    https://www.itl.nist.gov/div898/software/dataplot/refman2/auxillar/quantse.htm.
+    """
+    x = np.sort(x)
+    n = len(x)
+    m = int(q*n+0.5)
+    A = m - 1
+    B = n - m
+    i = np.arange(1,len(x)+1)
+    w = beta.cdf(i/n,A,B) - beta.cdf((i-1)/n,A,B)
+    return np.sqrt(np.sum(w*x**2)-np.sum(w*x)**2)
+
+def q90_err(x):
+    """
+    Returns the error on the 90th percentile for all the non NaN values in a
+    Series `x`.
+    """
+    x = x.dropna()
+    n = len(x)
+    if n == 0:
+        return np.nan
+    return quantile_error(x.values,0.9)
+
+def q90(x):
+    """
+    Returns the 90th percentile for all the non NaN values in a Series `x`.
+    """
+    x = x.dropna()
+    n = len(x)
+    if n == 0:
+        return np.nan
+    return np.percentile(x.values,90.0)
+
 def median(x):
     """
     Returns the median for all the non NaN values in a Series `x`.
@@ -243,13 +265,13 @@ def despine(fig=None, ax=None, top=True, right=True, left=False,
 
 if __name__ == '__main__':
     import argparse
-    import matplotlib.pyplot as plt
     import numpy as np
     import h5py
     import pandas as pd
 
     parser = argparse.ArgumentParser("plot fit results")
     parser.add_argument("filenames", nargs='+', help="input files")
+    parser.add_argument("--save", action="store_true", default=False, help="save plots")
     args = parser.parse_args()
 
     # Read in all the data.
@@ -322,9 +344,46 @@ if __name__ == '__main__':
 
     markers = itertools.cycle(('o', 'v')) 
 
-    # Default figure size. Currently set to my monitor width and height so that
-    # things are properly formatted
-    FIGSIZE = (13.78,7.48)
+    if args.save:
+        # default \textwidth for a fullpage article in Latex is 16.50764 cm.
+        # You can figure this out by compiling the following TeX document:
+        #
+        #    \documentclass{article}
+        #    \usepackage{fullpage}
+        #    \usepackage{layouts}
+        #    \begin{document}
+        #    textwidth in cm: \printinunitsof{cm}\prntlen{\textwidth}
+        #    \end{document}
+
+        width = 16.50764
+        width /= 2.54 # cm -> inches
+        # According to this page:
+        # http://www-personal.umich.edu/~jpboyd/eng403_chap2_tuftegospel.pdf,
+        # Tufte suggests an aspect ratio of 1.5 - 1.6.
+        height = width/1.5
+        FIGSIZE = (width,height)
+
+        import matplotlib.pyplot as plt
+
+        font = {'family':'serif', 'serif': ['computer modern roman']}
+        plt.rc('font',**font)
+
+        plt.rc('text', usetex=True)
+    else:
+        # on retina screens, the default plots are way too small
+        # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
+        # Qt5 will scale everything using the dpi in ~/.Xresources
+        import matplotlib
+        matplotlib.use("Qt5Agg")
+
+        import matplotlib.pyplot as plt
+
+        # Default figure size. Currently set to my monitor width and height so that
+        # things are properly formatted
+        FIGSIZE = (13.78,7.48)
+
+        # Make the defalt font bigger
+        plt.rc('font', size=22)
 
     fig3, ax3 = plt.subplots(3,1,figsize=FIGSIZE,num=3,sharex=True)
     fig4, ax4 = plt.subplots(3,1,figsize=FIGSIZE,num=4,sharex=True)
@@ -338,7 +397,7 @@ if __name__ == '__main__':
         dx = events.groupby(pd_bins)['dx'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
         dy = events.groupby(pd_bins)['dy'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
         dz = events.groupby(pd_bins)['dz'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
-        theta = events.groupby(pd_bins)['theta'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
+        theta = events.groupby(pd_bins)['theta'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err,q90,q90_err])
 
         label = 'Muon' if id == IDP_MU_MINUS else 'Electron'
 
@@ -361,7 +420,7 @@ if __name__ == '__main__':
         ax4[2].errorbar(T,dz['iqr_std'],yerr=dz['iqr_std_err'],fmt=marker,label=label)
 
         plt.figure(5,figsize=FIGSIZE)
-        plt.errorbar(T,theta['iqr_std'],yerr=theta['iqr_std_err'],fmt=marker,label=label)
+        plt.errorbar(T,theta['std'],yerr=theta['std_err'],fmt=marker,label=label)
 
         plt.figure(6,figsize=FIGSIZE)
         plt.scatter(events['Te'],events['ratio'],marker=marker,label=label)
@@ -370,7 +429,6 @@ if __name__ == '__main__':
     despine(fig,trim=True)
     plt.xlabel("Kinetic Energy (MeV)")
     plt.ylabel(r"Energy bias (\%)")
-    plt.title("Energy Bias")
     plt.legend()
     plt.tight_layout()
 
@@ -378,7 +436,6 @@ if __name__ == '__main__':
     despine(fig,trim=True)
     plt.xlabel("Kinetic Energy (MeV)")
     plt.ylabel(r"Energy resolution (\%)")
-    plt.title("Energy Resolution")
     plt.legend()
     plt.tight_layout()
 
@@ -392,7 +449,6 @@ if __name__ == '__main__':
     despine(ax=ax3[0],trim=True)
     despine(ax=ax3[1],trim=True)
     despine(ax=ax3[2],trim=True)
-    fig3.suptitle("Position Bias (cm)")
     h,l = ax3[0].get_legend_handles_labels()
     fig3.legend(h,l,loc='upper right')
     fig3.subplots_adjust(right=0.75)
@@ -409,7 +465,6 @@ if __name__ == '__main__':
     despine(ax=ax4[0],trim=True)
     despine(ax=ax4[1],trim=True)
     despine(ax=ax4[2],trim=True)
-    fig4.suptitle("Position Resolution (cm)")
     h,l = ax4[0].get_legend_handles_labels()
     fig4.legend(h,l,loc='upper right')
     fig4.subplots_adjust(right=0.75)
@@ -420,7 +475,6 @@ if __name__ == '__main__':
     despine(fig,trim=True)
     plt.xlabel("Kinetic Energy (MeV)")
     plt.ylabel("Angular resolution (deg)")
-    plt.title("Angular Resolution")
     plt.ylim((0,plt.gca().get_ylim()[1]))
     plt.legend()
     plt.tight_layout()
@@ -431,8 +485,40 @@ if __name__ == '__main__':
     despine(fig,trim=True)
     plt.xlabel("Reconstructed Electron Energy (MeV)")
     plt.ylabel(r"Log Likelihood Ratio (e/$\mu$)")
-    plt.title("Log Likelihood Ratio vs Reconstructed Electron Energy")
     plt.legend()
     plt.tight_layout()
 
-    plt.show()
+    if args.save:
+        fig = plt.figure(1)
+        plt.savefig("energy_bias.pdf")
+        plt.savefig("energy_bias.eps")
+        fig = plt.figure(2)
+        plt.savefig("energy_resolution.pdf")
+        plt.savefig("energy_resolution.eps")
+        fig = plt.figure(3)
+        plt.savefig("position_bias.pdf")
+        plt.savefig("position_bias.eps")
+        fig = plt.figure(4)
+        plt.savefig("position_resolution.pdf")
+        plt.savefig("position_resolution.eps")
+        fig = plt.figure(5)
+        plt.savefig("angular_resolution.pdf")
+        plt.savefig("angular_resolution.eps")
+        fig = plt.figure(6)
+        plt.savefig("likelihood_ratio.pdf")
+        plt.savefig("likelihood_ratio.eps")
+    else:
+        plt.figure(1)
+        plt.title("Energy Bias")
+        plt.figure(2)
+        plt.title("Energy Resolution")
+        plt.figure(3)
+        fig3.suptitle("Position Bias (cm)")
+        plt.figure(4)
+        fig4.suptitle("Position Resolution (cm)")
+        plt.figure(5)
+        plt.title("Angular Resolution")
+        plt.figure(6)
+        plt.title("Log Likelihood Ratio vs Reconstructed Electron Energy")
+
+        plt.show()