update script to calculate contamination

This commit updates the dc script to calculate the instrumental contamination
to now treat all 4 high level variables as correlated for muons. Previously I
had assumed that the reconstructed radius was independent from udotr, z, and
psi, but based on the corner plots it seems like the radius is strongly
correlated with udotr.

I also updated the plotting code when using the save command line argument to
be similar to plot-fit-results.

1 files changed, 187 insertions, 113 deletions

diff --git a/utils/dc b/utils/dc
index 30c5123..56decf9 100755
--- a/utils/dc
+++ b/utils/dc
@@ -33,25 +33,6 @@ except ImportError:
     print("emcee version 2.2.1 is required",file=sys.stderr)
     sys.exit(1)
 
-# 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")
-
-# Default figure size. Currently set to my monitor width and height so that
-# things are properly formatted
-FIGSIZE = (13.78,7.48)
-
-matplotlib.rcParams['figure.figsize'] = FIGSIZE
-
-matplotlib.rc('font', size=22)
-
-font = {'family':'serif', 'serif': ['computer modern roman']}
-matplotlib.rc('font',**font)
-
-matplotlib.rc('text', usetex=True)
-
 # Taken from https://raw.githubusercontent.com/mwaskom/seaborn/c73055b2a9d9830c6fbbace07127c370389d04dd/seaborn/utils.py
 def despine(fig=None, ax=None, top=True, right=True, left=False,
             bottom=False, offset=None, trim=False):
@@ -305,12 +286,6 @@ class bcolors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
-# 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")
-
 SNOMAN_MASS = {
     20: 0.511,
     21: 0.511,
@@ -591,35 +566,35 @@ class Constraint(object):
                 x[i] -= x[i]/2.0
         return x
 
-# Constraint to enforce the fact that P(psi,z,udotr|muon) all add up to 1.0. In
-# the likelihood function we set the last possibility for r and udotr equal to
-# 1.0 minus the others. Therefore, we need to enforce the fact that the others
-# must add up to less than 1.
-muon_psi_z_udotr = Constraint(range(12,19))
+# Constraint to enforce the fact that P(r,psi,z,udotr|muon) all add up to 1.0.
+# In the likelihood function we set the last possibility for r and udotr equal
+# to 1.0 minus the others. Therefore, we need to enforce the fact that the
+# others must add up to less than 1.
+muon_r_psi_z_udotr = Constraint(range(11,26))
 
 # Constraint to enforce the fact that P(z,udotr|noise) all add up to 1.0. In
 # the likelihood function we set the last possibility for r and udotr equal to
 # 1.0 minus the others. Therefore, we need to enforce the fact that the others
 # must add up to less than 1.
-noise_z_udotr = Constraint(range(21,24))
+noise_z_udotr = Constraint(range(28,31))
 
 # Constraint to enforce the fact that P(r,z,udotr|neck) all add up to 1.0. In
 # the likelihood function we set the last possibility for r and udotr equal to
 # 1.0 minus the others. Therefore, we need to enforce the fact that the others
 # must add up to less than 1.
-neck_r_z_udotr = Constraint(range(24,31))
+neck_r_z_udotr = Constraint(range(31,38))
 
 # Constraint to enforce the fact that P(r,udotr|flasher) all add up to 1.0. In
 # the likelihood function we set the last possibility for r and udotr equal to
 # 1.0 minus the others. Therefore, we need to enforce the fact that the others
 # must add up to less than 1
-flasher_r_udotr = Constraint(range(32,35))
+flasher_r_udotr = Constraint(range(39,42))
 
 # Constraint to enforce the fact that P(r,udotr|breakdown) all add up to 1.0.
 # In the likelihood function we set the last possibility for r and udotr equal
 # to 1.0 minus the others. Therefore, we need to enforce the fact that the
 # others must add up to less than 1.
-breakdown_r_udotr = Constraint(range(37,40))
+breakdown_r_udotr = Constraint(range(44,47))
 
 def make_nll(data, sacrifice,constraints):
     def nll(x, grad=None, fill_value=1e9):
@@ -641,47 +616,73 @@ def make_nll(data, sacrifice,constraints):
         if nll:
             return nll
 
-        mu_signal, mu_muon, mu_noise, mu_neck, mu_flasher, mu_breakdown, \
-        contamination_muon, contamination_noise, contamination_neck, contamination_flasher, contamination_breakdown, \
-        p_r_muon_lo, \
-        p_psi_z_udotr_muon_lololo, \
-        p_psi_z_udotr_muon_lolohi, \
-        p_psi_z_udotr_muon_lohilo, \
-        p_psi_z_udotr_muon_lohihi, \
-        p_psi_z_udotr_muon_hilolo, \
-        p_psi_z_udotr_muon_hilohi, \
-        p_psi_z_udotr_muon_hihilo, \
-        p_r_noise_lo, p_psi_noise_lo, \
-        p_z_udotr_noise_lolo, \
-        p_z_udotr_noise_lohi, \
-        p_z_udotr_noise_hilo, \
-        p_r_z_udotr_neck_lololo, \
-        p_r_z_udotr_neck_lolohi, \
-        p_r_z_udotr_neck_lohilo, \
-        p_r_z_udotr_neck_lohihi, \
-        p_r_z_udotr_neck_hilolo, \
-        p_r_z_udotr_neck_hilohi, \
-        p_r_z_udotr_neck_hihilo, \
-        p_psi_neck_lo, \
-        p_r_udotr_flasher_lolo, p_r_udotr_flasher_lohi, p_r_udotr_flasher_hilo, p_psi_flasher_lo, p_z_flasher_lo, \
-        p_r_udotr_breakdown_lolo, p_r_udotr_breakdown_lohi, p_r_udotr_breakdown_hilo, p_psi_breakdown_lo, p_z_breakdown_lo, \
-        p_neck_given_muon = x
+        (mu_signal, mu_muon, mu_noise, mu_neck, mu_flasher, mu_breakdown,
+         contamination_muon, contamination_noise, contamination_neck, contamination_flasher, contamination_breakdown,
+         p_r_psi_z_udotr_muon_lolololo, # 11
+         p_r_psi_z_udotr_muon_lololohi,
+         p_r_psi_z_udotr_muon_lolohilo,
+         p_r_psi_z_udotr_muon_lolohihi,
+         p_r_psi_z_udotr_muon_lohilolo,
+         p_r_psi_z_udotr_muon_lohilohi,
+         p_r_psi_z_udotr_muon_lohihilo,
+         p_r_psi_z_udotr_muon_lohihihi,
+         p_r_psi_z_udotr_muon_hilololo,
+         p_r_psi_z_udotr_muon_hilolohi,
+         p_r_psi_z_udotr_muon_hilohilo,
+         p_r_psi_z_udotr_muon_hilohihi,
+         p_r_psi_z_udotr_muon_hihilolo,
+         p_r_psi_z_udotr_muon_hihilohi,
+         p_r_psi_z_udotr_muon_hihihilo,
+         p_r_noise_lo, p_psi_noise_lo, # 26, 27
+         p_z_udotr_noise_lolo, # 28
+         p_z_udotr_noise_lohi,
+         p_z_udotr_noise_hilo,
+         p_r_z_udotr_neck_lololo, # 31
+         p_r_z_udotr_neck_lolohi,
+         p_r_z_udotr_neck_lohilo,
+         p_r_z_udotr_neck_lohihi,
+         p_r_z_udotr_neck_hilolo,
+         p_r_z_udotr_neck_hilohi,
+         p_r_z_udotr_neck_hihilo,
+         p_psi_neck_lo, # 38
+         p_r_udotr_flasher_lolo, p_r_udotr_flasher_lohi, p_r_udotr_flasher_hilo, # 39, ..., 41
+         p_psi_flasher_lo, p_z_flasher_lo,
+         p_r_udotr_breakdown_lolo, p_r_udotr_breakdown_lohi, p_r_udotr_breakdown_hilo, # 44, ..., 46
+         p_psi_breakdown_lo, p_z_breakdown_lo,
+         p_neck_given_muon) = x
 
         p_r_udotr_flasher_hihi = 1-p_r_udotr_flasher_lolo-p_r_udotr_flasher_lohi-p_r_udotr_flasher_hilo
         p_r_udotr_breakdown_hihi = 1-p_r_udotr_breakdown_lolo-p_r_udotr_breakdown_lohi-p_r_udotr_breakdown_hilo
-        p_psi_z_udotr_muon_hihihi = 1 - p_psi_z_udotr_muon_lololo - p_psi_z_udotr_muon_lolohi - p_psi_z_udotr_muon_lohilo - p_psi_z_udotr_muon_lohihi - p_psi_z_udotr_muon_hilolo - p_psi_z_udotr_muon_hilohi - p_psi_z_udotr_muon_hihilo
+        p_r_psi_z_udotr_muon_hihihihi = 1 - \
+            p_r_psi_z_udotr_muon_lolololo - \
+            p_r_psi_z_udotr_muon_lololohi - \
+            p_r_psi_z_udotr_muon_lolohilo - \
+            p_r_psi_z_udotr_muon_lolohihi - \
+            p_r_psi_z_udotr_muon_lohilolo - \
+            p_r_psi_z_udotr_muon_lohilohi - \
+            p_r_psi_z_udotr_muon_lohihilo - \
+            p_r_psi_z_udotr_muon_lohihihi - \
+            p_r_psi_z_udotr_muon_hilololo - \
+            p_r_psi_z_udotr_muon_hilolohi - \
+            p_r_psi_z_udotr_muon_hilohilo - \
+            p_r_psi_z_udotr_muon_hilohihi - \
+            p_r_psi_z_udotr_muon_hihilolo - \
+            p_r_psi_z_udotr_muon_hihilohi - \
+            p_r_psi_z_udotr_muon_hihihilo
         p_r_z_udotr_neck_hihihi = 1 - p_r_z_udotr_neck_lololo - p_r_z_udotr_neck_lolohi - p_r_z_udotr_neck_lohilo - p_r_z_udotr_neck_lohihi - p_r_z_udotr_neck_hilolo - p_r_z_udotr_neck_hilohi - p_r_z_udotr_neck_hihilo
         p_z_udotr_noise_hihi = 1 - p_z_udotr_noise_lolo - p_z_udotr_noise_lohi - p_z_udotr_noise_hilo
 
         # Muon events
         # first 6 parameters are the mean number of signal and bgs
-        p_r_muon = np.array([p_r_muon_lo,1-p_r_muon_lo])
-        p_psi_z_udotr_muon = np.array([\
-            [[p_psi_z_udotr_muon_lololo, p_psi_z_udotr_muon_lolohi], \
-             [p_psi_z_udotr_muon_lohilo, p_psi_z_udotr_muon_lohihi]], \
-            [[p_psi_z_udotr_muon_hilolo, p_psi_z_udotr_muon_hilohi], \
-             [p_psi_z_udotr_muon_hihilo, p_psi_z_udotr_muon_hihihi]]])
-        p_muon = p_r_muon[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_z_udotr_muon
+        p_muon = np.array([\
+            [[[p_r_psi_z_udotr_muon_lolololo, p_r_psi_z_udotr_muon_lololohi], \
+              [p_r_psi_z_udotr_muon_lolohilo, p_r_psi_z_udotr_muon_lolohihi]], \
+             [[p_r_psi_z_udotr_muon_lohilolo, p_r_psi_z_udotr_muon_lohilohi], \
+              [p_r_psi_z_udotr_muon_lohihilo, p_r_psi_z_udotr_muon_lohihihi]]], \
+            [[[p_r_psi_z_udotr_muon_hilololo, p_r_psi_z_udotr_muon_hilolohi], \
+              [p_r_psi_z_udotr_muon_hilohilo, p_r_psi_z_udotr_muon_hilohihi]], \
+             [[p_r_psi_z_udotr_muon_hihilolo, p_r_psi_z_udotr_muon_hihilohi], \
+              [p_r_psi_z_udotr_muon_hihihilo, p_r_psi_z_udotr_muon_hihihihi]]]])
         expected_muon = p_muon*contamination_muon*mu_muon + sacrifice['muon']*mu_signal
 
         nll -= poisson.logpmf(data['muon'],expected_muon).sum()
@@ -882,7 +883,6 @@ if __name__ == '__main__':
     # the *second* event after the breakdown didn't get tagged correctly. If we
     # apply the data cleaning cuts first and then tag prompt events then this
     # event will get tagged as a prompt event.
-    ev['prompt'] = (ev.nhit >= 100)
     ev = ev.groupby('run',as_index=False).apply(prompt_event).reset_index(level=0,drop=True)
 
     print("number of events after prompt nhit cut = %i" % np.count_nonzero(ev.prompt))
@@ -942,7 +942,7 @@ if __name__ == '__main__':
     p_udotr_signal = np.array([0.25,0.75])
     sacrifice['signal'] = p_r_signal[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_signal[:,np.newaxis,np.newaxis]*p_z_signal[:,np.newaxis]*p_udotr_signal
 
-    constraints = [flasher_r_udotr, breakdown_r_udotr,muon_psi_z_udotr,neck_r_z_udotr,noise_z_udotr]
+    constraints = [flasher_r_udotr, breakdown_r_udotr,muon_r_psi_z_udotr,neck_r_z_udotr,noise_z_udotr]
     nll = make_nll(data,sacrifice,constraints)
 
     x0 = []
@@ -953,18 +953,24 @@ if __name__ == '__main__':
     x0 += [0.99]*5
 
     if data['muon'].sum() > 0:
-        # P(r|muon)
-        x0 += [data['muon'][0].sum()/data['muon'].sum()]
-        # P(psi,z,udotr|muon)
-        x0 += [data['muon'][:,0,0,0].sum()/data['muon'].sum()]
-        x0 += [data['muon'][:,0,0,1].sum()/data['muon'].sum()]
-        x0 += [data['muon'][:,0,1,0].sum()/data['muon'].sum()]
-        x0 += [data['muon'][:,0,1,1].sum()/data['muon'].sum()]
-        x0 += [data['muon'][:,1,0,0].sum()/data['muon'].sum()]
-        x0 += [data['muon'][:,1,0,1].sum()/data['muon'].sum()]
-        x0 += [data['muon'][:,1,1,0].sum()/data['muon'].sum()]
+        # P(r,psi,z,udotr|muon)
+        x0 += [data['muon'][0,0,0,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,0,0,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,0,1,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,0,1,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,1,0,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,1,0,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,1,1,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][0,1,1,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,0,0,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,0,0,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,0,1,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,0,1,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,1,0,0].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,1,0,1].sum()/data['muon'].sum()]
+        x0 += [data['muon'][1,1,1,0].sum()/data['muon'].sum()]
     else:
-        x0 += [0.1]*8
+        x0 += [0.1]*15
 
     if data['noise'].sum() > 0:
         # P(r|noise)
@@ -992,7 +998,7 @@ if __name__ == '__main__':
     else:
         x0 += [0.1]*8
 
-    if data['neck'].sum() > 0:
+    if data['flasher'].sum() > 0:
         # P(r,udotr|flasher)
         x0 += [data['flasher'][0,:,:,0].sum()/data['flasher'].sum()]
         x0 += [data['flasher'][0,:,:,1].sum()/data['flasher'].sum()]
@@ -1078,13 +1084,54 @@ if __name__ == '__main__':
     print("Mean acceptance fraction: {0:.3f}".format(np.mean(sampler.acceptance_fraction)))
 
     try:
-        print("autocorrelation time: ", sampler.get_autocorr_time())
+        print("autocorrelation time: ", sampler.get_autocorr_time(quiet=True))
     except Exception as e:
         print(e)
 
+    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)
+
     # Plot walker positions as a function of step number for the expected
     # number of events
-    fig, axes = plt.subplots(6, num=1, sharex=True)
+    fig, axes = plt.subplots(6, figsize=FIGSIZE, num=1, sharex=True)
     samples = sampler.get_chain()
     labels = ["Signal","Muon","Noise","Neck","Flasher","Breakdown"]
     for i, bg in enumerate(['signal','muon','noise','neck','flasher','breakdown']):
@@ -1099,7 +1146,7 @@ if __name__ == '__main__':
 
     # Plot walker positions as a function of step number for the background cut
     # efficiencies
-    fig, axes = plt.subplots(5, num=2, sharex=True)
+    fig, axes = plt.subplots(5, figsize=FIGSIZE, num=2, sharex=True)
     samples = sampler.get_chain()
     tag_labels = ['M','N','Ne','F','B']
     for i, bg in enumerate(['muon','noise','neck','flasher','breakdown']):
@@ -1114,7 +1161,7 @@ if __name__ == '__main__':
 
     samples = sampler.chain.reshape((-1,len(x0)))
 
-    plt.figure(3)
+    plt.figure(3, figsize=FIGSIZE)
     for i, bg in enumerate(['signal','muon','noise','neck','flasher','breakdown']):
         ax = plt.subplot(3,2,i+1)
         plt.hist(samples[:,i],bins=100,histtype='step')
@@ -1124,7 +1171,7 @@ if __name__ == '__main__':
     plt.legend()
     plt.tight_layout()
 
-    plt.figure(4)
+    plt.figure(4, figsize=FIGSIZE)
     for i, bg in enumerate(['muon','noise','neck','flasher','breakdown']):
         ax = plt.subplot(3,2,i+1)
         plt.hist(samples[:,6+i],bins=100,histtype='step')
@@ -1134,31 +1181,58 @@ if __name__ == '__main__':
     plt.legend()
     plt.tight_layout()
 
-    mu_signal, mu_muon, mu_noise, mu_neck, mu_flasher, mu_breakdown, \
-    contamination_muon, contamination_noise, contamination_neck, contamination_flasher, contamination_breakdown, \
-    p_r_muon_lo, \
-    p_psi_z_udotr_muon_lololo, \
-    p_psi_z_udotr_muon_lolohi, \
-    p_psi_z_udotr_muon_lohilo, \
-    p_psi_z_udotr_muon_lohihi, \
-    p_psi_z_udotr_muon_hilolo, \
-    p_psi_z_udotr_muon_hilohi, \
-    p_psi_z_udotr_muon_hihilo, \
-    p_r_noise_lo, p_psi_noise_lo, \
-    p_z_udotr_noise_lolo, \
-    p_z_udotr_noise_lohi, \
-    p_z_udotr_noise_hilo, \
-    p_r_z_udotr_neck_lololo, \
-    p_r_z_udotr_neck_lolohi, \
-    p_r_z_udotr_neck_lohilo, \
-    p_r_z_udotr_neck_lohihi, \
-    p_r_z_udotr_neck_hilolo, \
-    p_r_z_udotr_neck_hilohi, \
-    p_r_z_udotr_neck_hihilo, \
-    p_psi_neck_lo, \
-    p_r_udotr_flasher_lolo, p_r_udotr_flasher_lohi, p_r_udotr_flasher_hilo, p_psi_flasher_lo, p_z_flasher_lo, \
-    p_r_udotr_breakdown_lolo, p_r_udotr_breakdown_lohi, p_r_udotr_breakdown_hilo, p_psi_breakdown_lo, p_z_breakdown_lo, \
-    p_neck_given_muon = samples.T
+    (mu_signal, mu_muon, mu_noise, mu_neck, mu_flasher, mu_breakdown,
+     contamination_muon, contamination_noise, contamination_neck, contamination_flasher, contamination_breakdown,
+     p_r_psi_z_udotr_muon_lolololo, # 11
+     p_r_psi_z_udotr_muon_lololohi,
+     p_r_psi_z_udotr_muon_lolohilo,
+     p_r_psi_z_udotr_muon_lolohihi,
+     p_r_psi_z_udotr_muon_lohilolo,
+     p_r_psi_z_udotr_muon_lohilohi,
+     p_r_psi_z_udotr_muon_lohihilo,
+     p_r_psi_z_udotr_muon_lohihihi,
+     p_r_psi_z_udotr_muon_hilololo,
+     p_r_psi_z_udotr_muon_hilolohi,
+     p_r_psi_z_udotr_muon_hilohilo,
+     p_r_psi_z_udotr_muon_hilohihi,
+     p_r_psi_z_udotr_muon_hihilolo,
+     p_r_psi_z_udotr_muon_hihilohi,
+     p_r_psi_z_udotr_muon_hihihilo,
+     p_r_noise_lo, p_psi_noise_lo, # 26, 27
+     p_z_udotr_noise_lolo, # 28
+     p_z_udotr_noise_lohi,
+     p_z_udotr_noise_hilo,
+     p_r_z_udotr_neck_lololo, # 31
+     p_r_z_udotr_neck_lolohi,
+     p_r_z_udotr_neck_lohilo,
+     p_r_z_udotr_neck_lohihi,
+     p_r_z_udotr_neck_hilolo,
+     p_r_z_udotr_neck_hilohi,
+     p_r_z_udotr_neck_hihilo,
+     p_psi_neck_lo, # 38
+     p_r_udotr_flasher_lolo, p_r_udotr_flasher_lohi, p_r_udotr_flasher_hilo, # 39, ..., 41
+     p_psi_flasher_lo, p_z_flasher_lo,
+     p_r_udotr_breakdown_lolo, p_r_udotr_breakdown_lohi, p_r_udotr_breakdown_hilo, # 44, ..., 46
+     p_psi_breakdown_lo, p_z_breakdown_lo,
+     p_neck_given_muon) = samples.T
+
+    p_r_muon_lo = p_r_psi_z_udotr_muon_lolololo + \
+                  p_r_psi_z_udotr_muon_lololohi + \
+                  p_r_psi_z_udotr_muon_lolohilo + \
+                  p_r_psi_z_udotr_muon_lolohihi + \
+                  p_r_psi_z_udotr_muon_lohilolo + \
+                  p_r_psi_z_udotr_muon_lohilohi + \
+                  p_r_psi_z_udotr_muon_lohihilo + \
+                  p_r_psi_z_udotr_muon_lohihihi
+
+    p_psi_muon_lo = p_r_psi_z_udotr_muon_lolololo + \
+                    p_r_psi_z_udotr_muon_lololohi + \
+                    p_r_psi_z_udotr_muon_lolohilo + \
+                    p_r_psi_z_udotr_muon_lolohihi + \
+                    p_r_psi_z_udotr_muon_hilololo + \
+                    p_r_psi_z_udotr_muon_hilolohi + \
+                    p_r_psi_z_udotr_muon_hilohilo + \
+                    p_r_psi_z_udotr_muon_hilohihi
 
     p_r = [sacrifice['signal'][0].sum(), p_r_muon_lo, p_r_noise_lo, \
         p_r_z_udotr_neck_lololo + p_r_z_udotr_neck_lolohi + p_r_z_udotr_neck_lohilo + p_r_z_udotr_neck_lohihi, \
@@ -1166,14 +1240,14 @@ if __name__ == '__main__':
         p_r_udotr_breakdown_lolo + p_r_udotr_breakdown_lohi]
 
     p_psi = [sacrifice['signal'][:,0].sum(), \
-        p_psi_z_udotr_muon_lololo + p_psi_z_udotr_muon_lolohi + p_psi_z_udotr_muon_lohilo + p_psi_z_udotr_muon_lohihi,
+        p_psi_muon_lo, \
         p_psi_noise_lo, \
         p_psi_neck_lo, \
         p_psi_flasher_lo, \
         p_psi_breakdown_lo]
 
     ylim_max = 0
-    fig = plt.figure(5)
+    fig = plt.figure(5, figsize=FIGSIZE)
     axes = []
     for i, bg in enumerate(['signal','muon','noise','neck','flasher','breakdown']):
         axes.append(plt.subplot(3,2,i+1))