update dc flasher and noise code

1 files changed, 45 insertions, 44 deletions

diff --git a/utils/dc b/utils/dc
index ee9a1c5..c6b2f36 100755
--- a/utils/dc
+++ b/utils/dc
@@ -62,11 +62,11 @@ def z_cut(ev):
 # 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
+# Constraint to enforce the fact that P(r,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(28,31))
+noise_r_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
@@ -74,11 +74,11 @@ noise_z_udotr = Constraint(range(28,31))
 # must add up to less than 1.
 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
+# Constraint to enforce the fact that P(z,udotr|flasher) all add up to 1.0. In
+# the likelihood function we set the last possibility for z 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(39,42))
+flasher_z_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
@@ -123,10 +123,10 @@ def make_nll(data, sacrifice, constraints, fitted_fraction):
          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_z_noise_lo, p_psi_noise_lo, # 26, 27
+         p_r_udotr_noise_lolo, # 28
+         p_r_udotr_noise_lohi,
+         p_r_udotr_noise_hilo,
          p_r_z_udotr_neck_lololo, # 31
          p_r_z_udotr_neck_lolohi,
          p_r_z_udotr_neck_lohilo,
@@ -135,13 +135,13 @@ def make_nll(data, sacrifice, constraints, fitted_fraction):
          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_z_udotr_flasher_lolo, p_z_udotr_flasher_lohi, p_z_udotr_flasher_hilo, # 39, ..., 41
+         p_r_flasher_lo, p_psi_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_z_udotr_flasher_hihi = 1 - p_z_udotr_flasher_lolo - p_z_udotr_flasher_lohi - p_z_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_r_psi_z_udotr_muon_hihihihi = 1 - \
             p_r_psi_z_udotr_muon_lolololo - \
@@ -160,7 +160,7 @@ def make_nll(data, sacrifice, constraints, fitted_fraction):
             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
+        p_r_udotr_noise_hihi = 1 - p_r_udotr_noise_lolo - p_r_udotr_noise_lohi - p_r_udotr_noise_hilo
 
         # Muon events
         # first 6 parameters are the mean number of signal and bgs
@@ -178,12 +178,12 @@ def make_nll(data, sacrifice, constraints, fitted_fraction):
         nll -= fast_poisson_logpmf(data['muon'],expected_muon).sum()
 
         # Noise events
-        p_r_noise = np.array([p_r_noise_lo,1-p_r_noise_lo])
+        p_r_udotr_noise = np.array([\
+            [p_r_udotr_noise_lolo,p_r_udotr_noise_lohi], \
+            [p_r_udotr_noise_hilo,p_r_udotr_noise_hihi]])
         p_psi_noise = np.array([p_psi_noise_lo,1-p_psi_noise_lo])
-        p_z_udotr_noise = np.array([\
-            [p_z_udotr_noise_lolo,p_z_udotr_noise_lohi],
-            [p_z_udotr_noise_hilo,p_z_udotr_noise_hihi]])
-        p_noise = p_r_noise[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_noise[:,np.newaxis,np.newaxis]*p_z_udotr_noise
+        p_z_noise = np.array([p_z_noise_lo,1-p_z_noise_lo])
+        p_noise = p_r_udotr_noise[:,np.newaxis,np.newaxis,:]*p_psi_noise[:,np.newaxis,np.newaxis]*p_z_noise[:,np.newaxis]
         expected_noise = p_noise*contamination_noise*mu_noise*fitted_fraction['noise'] + sacrifice['noise']*mu_signal
 
         nll -= fast_poisson_logpmf(data['noise'],expected_noise).sum()
@@ -204,12 +204,12 @@ def make_nll(data, sacrifice, constraints, fitted_fraction):
         nll -= fast_poisson_logpmf(data['neck'],expected_neck).sum()
 
         # Flasher events
-        p_r_udotr_flasher = np.array([\
-            [p_r_udotr_flasher_lolo,p_r_udotr_flasher_lohi], \
-            [p_r_udotr_flasher_hilo,p_r_udotr_flasher_hihi]])
+        p_r_flasher = np.array([p_r_flasher_lo,1-p_r_flasher_lo])
         p_psi_flasher = np.array([p_psi_flasher_lo,1-p_psi_flasher_lo])
-        p_z_flasher = np.array([p_z_flasher_lo,1-p_z_flasher_lo])
-        p_flasher = p_r_udotr_flasher[:,np.newaxis,np.newaxis,:]*p_psi_flasher[:,np.newaxis,np.newaxis]*p_z_flasher[:,np.newaxis]
+        p_z_udotr_flasher = np.array([\
+            [p_z_udotr_flasher_lolo,p_z_udotr_flasher_lohi],
+            [p_z_udotr_flasher_hilo,p_z_udotr_flasher_hihi]])
+        p_flasher = p_r_flasher[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_flasher[:,np.newaxis,np.newaxis]*p_z_udotr_flasher
         expected_flasher = p_flasher*contamination_flasher*mu_flasher*fitted_fraction['flasher'] + sacrifice['flasher']*mu_signal
 
         nll -= fast_poisson_logpmf(data['flasher'],expected_flasher).sum()
@@ -349,7 +349,7 @@ if __name__ == '__main__':
 
         sacrifice[bg] /= len(ev_mc)
 
-    constraints = [flasher_r_udotr, breakdown_r_udotr,muon_r_psi_z_udotr,neck_r_z_udotr,noise_z_udotr]
+    constraints = [flasher_z_udotr, breakdown_r_udotr,muon_r_psi_z_udotr,neck_r_z_udotr,noise_r_udotr]
     nll = make_nll(data,sacrifice,constraints,fitted_fraction)
 
     x0 = []
@@ -380,14 +380,14 @@ if __name__ == '__main__':
         x0 += [0.1]*15
 
     if data['noise'].sum() > 0:
-        # P(r|noise)
-        x0 += [data['noise'][0].sum()/data['noise'].sum()]
+        # P(z|noise)
+        x0 += [data['noise'][:,:,0].sum()/data['noise'].sum()]
         # P(psi|noise)
         x0 += [data['noise'][:,0].sum()/data['noise'].sum()]
-        # P(z,udotr|noise)
-        x0 += [data['noise'][:,:,0,0].sum()/data['noise'].sum()]
-        x0 += [data['noise'][:,:,0,1].sum()/data['noise'].sum()]
-        x0 += [data['noise'][:,:,1,0].sum()/data['noise'].sum()]
+        # P(r,udotr|noise)
+        x0 += [data['noise'][0,:,:,0].sum()/data['noise'].sum()]
+        x0 += [data['noise'][0,:,:,1].sum()/data['noise'].sum()]
+        x0 += [data['noise'][1,:,:,0].sum()/data['noise'].sum()]
     else:
         x0 += [0.1]*5
 
@@ -406,14 +406,14 @@ if __name__ == '__main__':
         x0 += [0.1]*8
 
     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()]
-        x0 += [data['flasher'][1,:,:,0].sum()/data['flasher'].sum()]
+        # P(z,udotr|flasher)
+        x0 += [data['flasher'][:,:,0,0].sum()/data['flasher'].sum()]
+        x0 += [data['flasher'][:,:,0,1].sum()/data['flasher'].sum()]
+        x0 += [data['flasher'][:,:,1,0].sum()/data['flasher'].sum()]
+        # P(r|flasher)
+        x0 += [data['flasher'][0].sum()/data['flasher'].sum()]
         # P(psi|flasher)
         x0 += [data['flasher'][:,0].sum()/data['flasher'].sum()]
-        # P(z|flasher)
-        x0 += [data['flasher'][:,:,0].sum()/data['flasher'].sum()]
     else:
         x0 += [0.1]*5
 
@@ -564,10 +564,10 @@ if __name__ == '__main__':
      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_z_noise_lo, p_psi_noise_lo, # 26, 27
+     p_r_udotr_noise_lolo, # 28
+     p_r_udotr_noise_lohi,
+     p_r_udotr_noise_hilo,
      p_r_z_udotr_neck_lololo, # 31
      p_r_z_udotr_neck_lolohi,
      p_r_z_udotr_neck_lohilo,
@@ -576,8 +576,8 @@ if __name__ == '__main__':
      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_z_udotr_flasher_lolo, p_z_udotr_flasher_lohi, p_z_udotr_flasher_hilo, # 39, ..., 41
+     p_r_flasher_lo, p_psi_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
@@ -600,9 +600,10 @@ if __name__ == '__main__':
                     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 = [sacrifice['signal'][0].sum(), p_r_muon_lo, \
+        p_r_udotr_noise_lolo + p_r_udotr_noise_lohi, \
         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_udotr_flasher_lolo + p_r_udotr_flasher_lohi, \
+        p_r_flasher_lo, \
         p_r_udotr_breakdown_lolo + p_r_udotr_breakdown_lohi]
 
     p_psi = [sacrifice['signal'][:,0].sum(), \