import numpy as np
from pax import plugin, utils
from pax.InterpolatingMap import InterpolatingMap
from scipy.special import betainc, gammaln
def bdtrc(k, n, p):
if (k < 0):
return (1.0)
if (k == n):
return (0.0)
dn = n - k
if (k == 0):
if (p < .01):
dk = -np.expm1(dn * np.log1p(-p))
else:
dk = 1.0 - np.exp(dn * np.log(1.0 - p))
else:
dk = k + 1
dk = betainc(dk, dn, p)
return dk
def bdtr(k, n, p):
if (k < 0):
return np.nan
if (k == n):
return (1.0)
dn = n - k
if (k == 0):
dk = np.exp(dn * np.log(1.0 - p))
else:
dk = k + 1
dk = betainc(dn, dk, 1.0 - p)
return dk
def binom_pmf(k, n, p):
scale_log = gammaln(n + 1) - gammaln(n - k + 1) - gammaln(k + 1)
ret_log = scale_log + k * np.log(p) + (n - k) * np.log(1 - p)
return np.exp(ret_log)
def binom_cdf(k, n, p):
return bdtr(k, n, p)
def binom_sf(k, n, p):
return bdtrc(k, n, p)
def binom_test(k, n, p):
'''
The main purpose of this algorithm is to find the value j on the
other side of the mean that has the same probability as k, and
integrate the tails outward from k and j. In the case where either
k or j are zero, only the non-zero tail is integrated.
'''
if n < k:
raise ValueError("n must be >= k")
if (p > 1.0) or (p < 0.0):
raise ValueError("p must be in range [0, 1]")
if k < 0:
raise ValueError("k must be >= 0")
d = binom_pmf(k, n, p)
rerr = 1 + 1e-7
d = d * rerr
n_iter = int(max(np.round(np.log10(n)) + 1, 2))
if k < n * p:
j_min, j_max = k, n
def check(d, y0, y1):
return ((y0 >= d) and (d > y1))
else:
if binom_pmf(0, n, p) > d:
j_min = j_max = 0
n_iter = 0
else:
j_min, j_max = 0, k
def check(d, y0, y1):
return ((y0 <= d) and (d < y1))
for i in range(n_iter): # successive approximation loop
n_pts = 20 if i == 0 else 10
j_range = np.linspace(j_min, j_max, n_pts, endpoint=True)
y = binom_pmf(j_range, n, p)
for i in range(len(j_range) - 1):
if check(d, y[i], y[i + 1]):
j_min, j_max = j_range[i], j_range[i + 1]
break
j = max(min((j_min + j_max) / 2, n), 0)
if k * j == 0: # one is zero, means we do a one-sided test
pval = binom_sf(max(k, j), n, p)
else:
pval = binom_cdf(min(k, j), n, p) + binom_sf(max(k, j), n, p)
return min(1.0, pval)
def s1_area_fraction_top_probability(aft_prob, area_tot, area_fraction_top,
hits_tot, hits_fraction_top, low_pe_threshold=10,
mode='test'):
"""Wrapper that does the S1 AFT probability calculation for you
"""
# below this in PE, transition to hits
if area_tot < low_pe_threshold:
s1_frac = area_tot / low_pe_threshold
hits_top = hits_tot * hits_fraction_top
s1_top = area_tot * area_fraction_top
size_top = hits_top * (1. - s1_frac) + s1_top * s1_frac
size_tot = hits_tot * (1. - s1_frac) + area_tot * s1_frac
else:
size_top = area_tot * area_fraction_top
size_tot = area_tot
if mode == 'pmf':
return binom_pmf(size_top, size_tot, aft_prob)
else:
return binom_test(size_top, size_tot, aft_prob)
class S1AreaFractionTopProbability(plugin.TransformPlugin):
"""Computes p-value for S1 area fraction top for each interaction
"""
def startup(self):
aftmap_filename = utils.data_file_name('XENON1T_s1_aft_xyz_20170808.json')
self.aft_map = InterpolatingMap(aftmap_filename)
def transform_event(self, event):
for ia in event.interactions:
s1 = event.peaks[ia.s1]
aft_prob = self.aft_map.get_value(ia.x, ia.y, ia.z)
ia.s1_area_fraction_top_probability = s1_area_fraction_top_probability(
aft_prob, s1.area, s1.area_fraction_top, s1.n_hits, s1.hits_fraction_top)
return event
