update plot-energy to calculate likelihood ratio

2 files changed, 23 insertions, 12 deletions

diff --git a/utils/plot-energy b/utils/plot-energy
index 6997f3b..ff36a19 100755
--- a/utils/plot-energy
+++ b/utils/plot-energy
@@ -104,7 +104,7 @@ if __name__ == '__main__':
     import sys
     import h5py
     from sddm.plot_energy import *
-    from sddm.plot import despine
+    from sddm.plot import *
     from sddm import setup_matplotlib
 
     parser = argparse.ArgumentParser("plot fit results")
diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py
index 4457dc0..62f6f34 100755
--- a/utils/sddm/plot_energy.py
+++ b/utils/sddm/plot_energy.py
@@ -16,11 +16,12 @@
 
 from __future__ import print_function, division
 import numpy as np
-from scipy.stats import poisson
+from scipy.stats import poisson, dirichlet, multinomial, beta
 from scipy.special import spence
 from .dc import DC_MUON, DC_JUNK, DC_CRATE_ISOTROPY, DC_QVNHIT, DC_NECK, DC_FLASHER, DC_ESUM, DC_OWL, DC_OWL_TRIGGER, DC_FTS, DC_ITC, DC_BREAKDOWN
 from . import grouper, print_warning, AV_RADIUS, PSUP_RADIUS
 import pandas as pd
+from scipy.misc import logsumexp
 
 # Fermi constant
 GF                      = 1.16637887e-5 # 1/MeV^2
@@ -198,6 +199,25 @@ def prompt_event(ev):
     ev.loc[ev.prompt,'prompt'] &= np.concatenate(([True],np.diff(ev[ev.prompt].gtr.values) > 250e6))
     return ev
 
+def sample_posterior_correlated(n,size):
+    # Assume prior is a dirichlet with alpha = 1
+    return dirichlet.rvs(n+1,size=size)
+
+def sample_posterior_independent(n,size):
+    p1 = beta.rvs(n[0]+n[1]+1,n[2]+n[3]+1,size=size)
+    p2 = beta.rvs(n[0]+n[2]+1,n[1]+n[3]+1,size=size)
+    return np.vstack((p1*p2,p1*(1-p2),(1-p1)*p2,(1-p1)*(1-p2))).T
+
+def nll(p,n):
+    return -multinomial.logpmf(n,np.sum(n),p)
+
+def get_likelihood_ratio(n,size=100000):
+    samples_correlated = sample_posterior_correlated(n,size)
+    nll_correlated = nll(samples_correlated,n)
+    samples_independent = sample_posterior_independent(n,size)
+    nll_independent = nll(samples_independent,n)
+    return logsumexp(-nll_independent) - logsumexp(-nll_correlated)
+
 def plot_corner_plot(ev, title, save=None):
     import matplotlib.pyplot as plt
     from .plot import despine
@@ -226,21 +246,12 @@ def plot_corner_plot(ev, title, save=None):
                 plt.gca().set_ylim(limits[i])
                 n = len(ev)
                 if n:
-                    p_i_lo = np.count_nonzero(ev[variables[i]] < cuts[i])/n
-                    p_j_lo = np.count_nonzero(ev[variables[j]] < cuts[j])/n
-                    p_lolo = p_i_lo*p_j_lo
-                    p_lohi = p_i_lo*(1-p_j_lo)
-                    p_hilo = (1-p_i_lo)*p_j_lo
-                    p_hihi = (1-p_i_lo)*(1-p_j_lo)
                     n_lolo = np.count_nonzero((ev[variables[i]] < cuts[i]) & (ev[variables[j]] < cuts[j]))
                     n_lohi = np.count_nonzero((ev[variables[i]] < cuts[i]) & (ev[variables[j]] >= cuts[j]))
                     n_hilo = np.count_nonzero((ev[variables[i]] >= cuts[i]) & (ev[variables[j]] < cuts[j]))
                     n_hihi = np.count_nonzero((ev[variables[i]] >= cuts[i]) & (ev[variables[j]] >= cuts[j]))
                     observed = np.array([n_lolo,n_lohi,n_hilo,n_hihi])
-                    expected = n*np.array([p_lolo,p_lohi,p_hilo,p_hihi])
-                    psi = -poisson.logpmf(observed,expected).sum() + poisson.logpmf(observed,observed).sum()
-                    psi /= np.std(-poisson.logpmf(np.random.poisson(observed,size=(10000,4)),observed).sum(axis=1) + poisson.logpmf(observed,observed).sum())
-                    plt.title(r"$\psi = %.1f$" % psi)
+                    plt.title(r"$\lambda = %.1f$" % get_likelihood_ratio(observed))
             if i == len(variables) - 1:
                 plt.xlabel(labels[j])
             else: