In [6]:
import numpy as np
import matplotlib.pyplot as plt
import scipy.special as spsp
In [7]:
def start(lmbda):
start_x = int(lmbda)
pmf = np.exp(-lmbda)
cdf = pmf
for x in range(1, start_x+1):
pmf = pmf * lmbda / x
cdf = cdf + pmf
return cdf, pmf
In [8]:
def poisson_sampling(lmbda, cdf_start, pmf_start):
x = int(lmbda)
cdf = cdf_start
pmf = pmf_start
u = np.random.rand()
if u >= cdf:
x = x + 1
pmf = pmf * lmbda / x
cdf = cdf + pmf
while u >= cdf:
x = x + 1
pmf = pmf * lmbda/x
cdf = cdf + pmf
return x
else:
cdf = cdf – pmf
while u < cdf:
pmf = pmf * x / lmbda
cdf = cdf - pmf
x = x - 1
return x
In [9]:
# Ordinarily, we would want a distribution with the same support
# as the target. However, the target distribution is just a Poisson
# variable, only truncated to values less than or equal to k. We can
# just use a Poisson variable as our proposal, then, since we keep all
# values <= k, and reject values > k.
def trunc_poisson(lmbda, k, n=1000):
res = []
cdf_start, pmf_start = start(lmbda)
for i in range(n):
x = poisson_sampling(lmbda, cdf_start, pmf_start)
while (x > k):
x = poisson_sampling(lmbda, cdf_start, pmf_start)
res.append(x)
return np.array(res)
# If you don’t want to do it that way, you can still use traditional
# rejection sampling!
def rejection(lmbda, k, n=1000):
pmf_proposal = np.ones(k+1) * 1/(k+1)
pmf_target = np.array([np.exp(lmbda) * lmbda**x / (spsp.factorial(x)) for x in range(k+1)])
pmf_target /= np.sum(pmf_target)
c = np.max(pmf_target / pmf_proposal)
res = []
for i in range(n):
x = int(np.random.rand() * (k + 1))
while np.random.rand() >= pmf_target[x] / pmf_proposal[x] / c:
x = int(np.random.rand() * (k + 1))
res.append(x)
return np.array(res)
In [10]:
pois_trunc = trunc_poisson(10,5)
print(np.mean(pois_trunc))
rejection_samples = rejection(10, 5)
print(np.mean(rejection_samples))
4.353
4.372