simplify nllr() function

1 files changed, 5 insertions, 9 deletions

diff --git a/utils/sddm/stats.py b/utils/sddm/stats.py
index 5324eea..ff2a6ad 100644
--- a/utils/sddm/stats.py
+++ b/utils/sddm/stats.py
@@ -8,7 +8,7 @@ particle_id = {20: 'e', 22: r'\mu'}
 def chi2(samples,expected):
     return np.sum((samples-expected)**2/expected,axis=-1)
 
-def nllr(samples,expected):
+def nllr(oi,ei):
     """
     Returns 2 times the negative log likelihood ratio that the `samples`
     histogram was drawn from the histogram `expected`. The numerator in the
@@ -22,14 +22,10 @@ def nllr(samples,expected):
     This likelihood ratio is supposed to give better results as a test
     statistic when computing p-values for histograms with low statistics.
     """
-    samples = np.atleast_2d(samples)
-    n = samples.sum(axis=-1)
-    N = expected.sum()
-    p_samples = samples.astype(np.double) + 1e-10
-    p_samples /= p_samples.sum(-1)[:,np.newaxis]
-    p_expected = expected.astype(np.double) + 1e-10
-    p_expected /= p_expected.sum()
-    return 2*(N - n + (samples*np.log(np.where(samples == 0,1,samples)/expected)).sum(axis=-1)).squeeze()[()]
+    oi = np.atleast_2d(oi)
+    n = oi.sum(axis=-1)
+    N = ei.sum()
+    return 2*(N - n + (oi*np.log(np.where(oi == 0,1,oi)/ei)).sum(axis=-1)).squeeze()[()]
 
 def get_mc_hist_posterior(hist, data, norm):
     """