Statistical numerics are a whole thing, and as a purist I would say that CodeReview does not include in its purview analysis of whether your code is correct. You need unit tests, including ones where you have chosen sample values where you've worked out - by hand, with a calculator - obviously-correct expectations. The unit tests I will show are "regression-only": they show that my suggested changes don't affect the output of your code, but don't say whether that output is correct.
If you're really not sure whether your math is correct (separable from the question of whether your code is well-written), consider instead a place like Cross-Validated.
Generally:
- Add type hints, including for your
alternative which - if it's not an Enum - should at least be a Literal
- Don't local-import; import at the top of the file
- Your
a/b calculations can be simplified
n, a and b can be put into a dictionary for kwargs parameter reuse- Don't pass
loc=0; that's already the default
- Rather than separate
x == p * n and x < p * n checks, consider just calculating a delta and comparing that to zero as needed
- Only calculate your
y in one place
- In the
x == p * n case, early-return to avoid a min call
- I'm not sure why you gave
n a default of None. I can't find a scenario where that would be valid.
- For obscure reasons,
count_nonzero is faster than sum
Suggested
from numbers import Real
from typing import Literal
from scipy import stats
import numpy as np
def beta_binom_test(
x: Real,
n: int,
p: Real = 0.5,
rho: Real = 0.1,
alternative: Literal['less', 'greater', 'two-sided'] = 'two-sided',
) -> float:
"""Using rho to find the shape parameters for the beta binomial
rho is empirically estimated"""
coefficient = 1/rho - 1
a = coefficient * p
params = {
'n': n,
'a': a, # shape parameter 1
'b': coefficient - a, # shape parameter 2
}
if alternative == 'less':
return stats.betabinom.cdf(k=x, **params)
if alternative == 'greater':
return stats.betabinom.sf(k=x - 1, **params)
# if alternative was neither 'less' nor 'greater', then it's 'two-sided'
delta = x - p*n
if delta == 0:
# special case as shortcut, would also be handled by `else` below
return 1
if delta < 0:
i = np.arange(np.ceil(p * n), n + 1)
else:
i = np.arange(np.floor(p * n) + 1)
d = stats.betabinom.pmf(k=x, **params)
rerr = 1 + 1e-7
y = np.count_nonzero(stats.betabinom.pmf(i, **params) <= d * rerr)
if delta < 0:
k_cdf = x
k_sf = n - y
else:
k_cdf = y - 1
k_sf = x - 1
pval = (stats.betabinom.cdf(k=k_cdf, **params) +
stats.betabinom.sf(k=k_sf, **params))
return min(1, pval)
def test() -> None:
def isclose(expected: Real, actual: Real) -> None:
assert np.isclose(expected, actual, rtol=0, atol=1e-12)
isclose(0, beta_binom_test(x=-1.3, n=10, p=0.5, rho=0.1, alternative='less'))
isclose(0.0125026702880859, beta_binom_test(x=0.3, n=10, p=0.5, rho=0.1, alternative='less'))
isclose(0.1000000000000000, beta_binom_test(x=0.3, n= 1, p=0.9, rho=0.2, alternative='less'))
isclose(0.5189135546558704, beta_binom_test(x=0.6, n= 5, p=0.2, rho=0.3, alternative='less'))
isclose(0.7289561524195800, beta_binom_test(x=0.5, n= 7, p=0.1, rho=0.4, alternative='less'))
isclose(0.7200000000000000, beta_binom_test(x=0.2, n= 2, p=0.2, rho=0.5, alternative='less'))
isclose(1, beta_binom_test(x=0.3, n=10, p=0.5, rho=0.4, alternative='greater'))
isclose(0.8771254595588235, beta_binom_test(x=1.3, n=10, p=0.5, rho=0.4, alternative='greater'))
isclose(0.2000000000000000, beta_binom_test(x=1.9, n= 1, p=0.2, rho=0.9, alternative='greater'))
isclose(0.4810864453441296, beta_binom_test(x=1.6, n= 5, p=0.2, rho=0.3, alternative='greater'))
isclose(0.2710438475804200, beta_binom_test(x=1.5, n= 7, p=0.1, rho=0.4, alternative='greater'))
isclose(0.2800000000000000, beta_binom_test(x=1.2, n= 2, p=0.2, rho=0.5, alternative='greater'))
# x < p*n
isclose(0.0125026702880859, beta_binom_test(x=0.3, n=10, p=0.5, rho=0.1, alternative='two-sided'))
isclose(0.1000000000000000, beta_binom_test(x=0.3, n= 1, p=0.9, rho=0.2, alternative='two-sided'))
isclose(0.5189135546558704, beta_binom_test(x=0.6, n= 5, p=0.2, rho=0.3, alternative='two-sided'))
isclose(0.7289561524195800, beta_binom_test(x=0.5, n= 7, p=0.1, rho=0.4, alternative='two-sided'))
isclose(0.7200000000000000, beta_binom_test(x=0.2, n= 2, p=0.2, rho=0.5, alternative='two-sided'))
# x == p*n
isclose(1, beta_binom_test(x=0.8, p=0.2, n=4, rho=0.5, alternative='two-sided'))
# x > p*n
isclose(1, beta_binom_test(x=0.9, p=0.2, n=4, rho=0.5, alternative='two-sided'))
if __name__ == '__main__':
test()
