From 40d962936761cb356d021fa1d4413974c8a65c7a Mon Sep 17 00:00:00 2001
From: tlatorre <tlatorre@uchicago.edu>
Date: Tue, 12 May 2020 14:02:58 -0500
Subject: speed up the contamination analysis script

---
 utils/dc         | 12 ++++++------
 utils/sddm/dc.py |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

(limited to 'utils')

diff --git a/utils/dc b/utils/dc
index 16591df..834c6e5 100755
--- a/utils/dc
+++ b/utils/dc
@@ -185,7 +185,7 @@ def make_nll(data, sacrifice, constraints):
               [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()
+        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])
@@ -196,7 +196,7 @@ def make_nll(data, sacrifice, constraints):
         p_noise = p_r_noise[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_noise[:,np.newaxis,np.newaxis]*p_z_udotr_noise
         expected_noise = p_noise*contamination_noise*mu_noise + sacrifice['noise']*mu_signal
 
-        nll -= poisson.logpmf(data['noise'],expected_noise).sum()
+        nll -= fast_poisson_logpmf(data['noise'],expected_noise).sum()
 
         # Neck events
         # FIXME: for now assume parameterized same as muon
@@ -211,7 +211,7 @@ def make_nll(data, sacrifice, constraints):
         # FIXME: pdf should be different for muon given neck
         expected_neck += p_muon*p_neck_given_muon*mu_muon
 
-        nll -= poisson.logpmf(data['neck'],expected_neck).sum()
+        nll -= fast_poisson_logpmf(data['neck'],expected_neck).sum()
 
         # Flasher events
         p_r_udotr_flasher = np.array([\
@@ -222,7 +222,7 @@ def make_nll(data, sacrifice, constraints):
         p_flasher = p_r_udotr_flasher[:,np.newaxis,np.newaxis,:]*p_psi_flasher[:,np.newaxis,np.newaxis]*p_z_flasher[:,np.newaxis]
         expected_flasher = p_flasher*contamination_flasher*mu_flasher + sacrifice['flasher']*mu_signal
 
-        nll -= poisson.logpmf(data['flasher'],expected_flasher).sum()
+        nll -= fast_poisson_logpmf(data['flasher'],expected_flasher).sum()
 
         # Breakdown events
         p_r_udotr_breakdown = np.array([\
@@ -233,7 +233,7 @@ def make_nll(data, sacrifice, constraints):
         p_breakdown = p_r_udotr_breakdown[:,np.newaxis,np.newaxis,:]*p_psi_breakdown[:,np.newaxis,np.newaxis]*p_z_breakdown[:,np.newaxis]
         expected_breakdown = p_breakdown*contamination_breakdown*mu_breakdown + sacrifice['breakdown']*mu_signal
 
-        nll -= poisson.logpmf(data['breakdown'],expected_breakdown).sum()
+        nll -= fast_poisson_logpmf(data['breakdown'],expected_breakdown).sum()
 
         # Signal like events
         expected_signal = np.zeros_like(expected_muon)
@@ -244,7 +244,7 @@ def make_nll(data, sacrifice, constraints):
         expected_signal += p_flasher*(1-contamination_flasher)*mu_flasher
         expected_signal += p_breakdown*(1-contamination_breakdown)*mu_breakdown
 
-        nll -= poisson.logpmf(data['signal'],expected_signal).sum()
+        nll -= fast_poisson_logpmf(data['signal'],expected_signal).sum()
 
         if not np.isfinite(nll):
             print("x = ", x)
diff --git a/utils/sddm/dc.py b/utils/sddm/dc.py
index 9081a21..a1d87b5 100755
--- a/utils/sddm/dc.py
+++ b/utils/sddm/dc.py
@@ -335,7 +335,7 @@ def prompt_event(ev):
 
 class Constraint(object):
     def __init__(self, range):
-        self.range = range
+        self.range = np.array(range)
 
     def __call__(self, x, grad=None):
         return np.sum(x[self.range]) - 1 + EPSILON
-- 
cgit