-
Notifications
You must be signed in to change notification settings - Fork 0
/
beta_fit.py
39 lines (31 loc) · 1.3 KB
/
beta_fit.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from typing import Optional, Tuple
import numpy as np
def naive_beta_binomial_fit(
data: list, n: Optional[int] = None
) -> Tuple[float, float, int]:
"""Estimate Beta-binomial distribution parameters from the mean and variance of the data.
Input:
data:
List (array) containing data values. Logic dictates that it should
contain positive integer values, as otherwise fitting Beta-binomial
distribution makes no sense.
n: (default: None)
N parameter of the Beta-binomial distribution can be fixed by
passing this optional value. If `None` is passed (which is the
default), then the maximum of the `data` will be used as the
estimate for N parameter.
Output:
alpha, beta and N parameter estimates.
"""
if n is None:
n = np.max(data)
mean = np.mean(data)
variance = np.var(data, ddof=1)
# these formulas were obtained by inverting the expressions for mean and
# variance of the Beta-binomial distribution (e.g., see
# https://en.wikipedia.org/wiki/Beta-binomial_distribution)
alpha_par = ((mean**2) * n - mean**3 - mean * variance) / (
mean**2 - mean * n + n * variance
)
beta_par = alpha_par * (n / mean - 1)
return alpha_par, beta_par, n