Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/statsmodels/distributions/discrete.py : 54%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import numpy as np
2from scipy.stats import rv_discrete, nbinom
3from scipy.special import gammaln
4from statsmodels.compat.scipy import _lazywhere
7class genpoisson_p_gen(rv_discrete):
8 '''Generalized Poisson distribution
9 '''
10 def _argcheck(self, mu, alpha, p):
11 return (mu >= 0) & (alpha==alpha) & (p > 0)
13 def _logpmf(self, x, mu, alpha, p):
14 mu_p = mu ** (p - 1.)
15 a1 = np.maximum(np.nextafter(0, 1), 1 + alpha * mu_p)
16 a2 = np.maximum(np.nextafter(0, 1), mu + (a1 - 1.) * x)
17 logpmf_ = np.log(mu) + (x - 1.) * np.log(a2)
18 logpmf_ -= x * np.log(a1) + gammaln(x + 1.) + a2 / a1
19 return logpmf_
21 def _pmf(self, x, mu, alpha, p):
22 return np.exp(self._logpmf(x, mu, alpha, p))
24genpoisson_p = genpoisson_p_gen(name='genpoisson_p',
25 longname='Generalized Poisson')
27class zipoisson_gen(rv_discrete):
28 '''Zero Inflated Poisson distribution
29 '''
30 def _argcheck(self, mu, w):
31 return (mu > 0) & (w >= 0) & (w<=1)
33 def _logpmf(self, x, mu, w):
34 return _lazywhere(x != 0, (x, mu, w),
35 (lambda x, mu, w: np.log(1. - w) + x * np.log(mu) -
36 gammaln(x + 1.) - mu),
37 np.log(w + (1. - w) * np.exp(-mu)))
39 def _pmf(self, x, mu, w):
40 return np.exp(self._logpmf(x, mu, w))
42zipoisson = zipoisson_gen(name='zipoisson',
43 longname='Zero Inflated Poisson')
45class zigeneralizedpoisson_gen(rv_discrete):
46 '''Zero Inflated Generalized Poisson distribution
47 '''
48 def _argcheck(self, mu, alpha, p, w):
49 return (mu > 0) & (w >= 0) & (w<=1)
51 def _logpmf(self, x, mu, alpha, p, w):
52 return _lazywhere(x != 0, (x, mu, alpha, p, w),
53 (lambda x, mu, alpha, p, w: np.log(1. - w) +
54 genpoisson_p.logpmf(x, mu, alpha, p)),
55 np.log(w + (1. - w) *
56 genpoisson_p.pmf(x, mu, alpha, p)))
58 def _pmf(self, x, mu, alpha, p, w):
59 return np.exp(self._logpmf(x, mu, alpha, p, w))
61zigenpoisson = zigeneralizedpoisson_gen(name='zigenpoisson',
62 longname='Zero Inflated Generalized Poisson')
64class zinegativebinomial_gen(rv_discrete):
65 '''Zero Inflated Generalized Negative Binomial distribution
66 '''
67 def _argcheck(self, mu, alpha, p, w):
68 return (mu > 0) & (w >= 0) & (w<=1)
70 def _logpmf(self, x, mu, alpha, p, w):
71 s, p = self.convert_params(mu, alpha, p)
72 return _lazywhere(x != 0, (x, s, p, w),
73 (lambda x, s, p, w: np.log(1. - w) +
74 nbinom.logpmf(x, s, p)),
75 np.log(w + (1. - w) *
76 nbinom.pmf(x, s, p)))
78 def _pmf(self, x, mu, alpha, p, w):
79 return np.exp(self._logpmf(x, mu, alpha, p, w))
81 def convert_params(self, mu, alpha, p):
82 size = 1. / alpha * mu**(2-p)
83 prob = size / (size + mu)
84 return (size, prob)
86zinegbin = zinegativebinomial_gen(name='zinegbin',
87 longname='Zero Inflated Generalized Negative Binomial')