Hide keyboard shortcuts

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 

5 

6 

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) 

12 

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_ 

20 

21 def _pmf(self, x, mu, alpha, p): 

22 return np.exp(self._logpmf(x, mu, alpha, p)) 

23 

24genpoisson_p = genpoisson_p_gen(name='genpoisson_p', 

25 longname='Generalized Poisson') 

26 

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) 

32 

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))) 

38 

39 def _pmf(self, x, mu, w): 

40 return np.exp(self._logpmf(x, mu, w)) 

41 

42zipoisson = zipoisson_gen(name='zipoisson', 

43 longname='Zero Inflated Poisson') 

44 

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) 

50 

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))) 

57 

58 def _pmf(self, x, mu, alpha, p, w): 

59 return np.exp(self._logpmf(x, mu, alpha, p, w)) 

60 

61zigenpoisson = zigeneralizedpoisson_gen(name='zigenpoisson', 

62 longname='Zero Inflated Generalized Poisson') 

63 

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) 

69 

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))) 

77 

78 def _pmf(self, x, mu, alpha, p, w): 

79 return np.exp(self._logpmf(x, mu, alpha, p, w)) 

80 

81 def convert_params(self, mu, alpha, p): 

82 size = 1. / alpha * mu**(2-p) 

83 prob = size / (size + mu) 

84 return (size, prob) 

85 

86zinegbin = zinegativebinomial_gen(name='zinegbin', 

87 longname='Zero Inflated Generalized Negative Binomial')