diff options
Diffstat (limited to 'rumba/recpoisson.py')
-rw-r--r-- | rumba/recpoisson.py | 65 |
1 files changed, 65 insertions, 0 deletions
diff --git a/rumba/recpoisson.py b/rumba/recpoisson.py new file mode 100644 index 0000000..3c1e6fe --- /dev/null +++ b/rumba/recpoisson.py @@ -0,0 +1,65 @@ +import math +import random + +import sys + +if sys.version_info < (3, 2): + from repoze.lru import lru_cache + # from functools32 import lru_cache +else: + from functools import lru_cache + + +@lru_cache(1000) +def _get_poisson_var(parameter): + return Poisson(parameter) + + +class Poisson(object): + + def __init__(self, parameter): + self.parameter = parameter + + def c_p(k): + """Compute the Poisson CDF for k iteratively.""" + if k == 0: + return self._p(0) + else: + return self._compute_poisson_cdf(k - 1) + self._p(k) + self._compute_poisson_cdf = lru_cache(int(2.5*self.parameter) + 1)(c_p) + + @staticmethod + def _get_random(): + return random.random() + + def _p(self, k): + # l^k * e^-l / k! + # Computed as exp(klog(l) - l - log(k!)) + l = self.parameter + l_to_the_k = k * math.log(l) + k_fact = sum([math.log(i + 1) for i in range(k)]) + return math.exp(l_to_the_k - l - k_fact) + + def sample(self): + # The idea is: + # take a sample from U(0,1) and call it f. + # Let x be s.t. x = min_N F(x) > f + # where F is the cumulative distribution function of Poisson(parameter) + # return x + f = self._get_random() + + # We compute x iteratively by computing + # \sum_k(P=k) + # where P ~ Poisson(parameter) and stopping as soon as + # it is greater than f. + # We use the cache to store results. + current_cdf = -1 + current_x = -1 + while current_cdf < f: + current_x += 1 + current_cdf = self._compute_poisson_cdf(current_x) + return current_x + + +def poisson(parameter): + return _get_poisson_var(parameter).sample() |