Scipy poisson distribution with an upper limit

Question

I am generating a random number using scipy stats. I used the Poisson distribution. Below is an example:

import scipy.stats as sct

A =2.5
Pos = sct.poisson.rvs(A,size = 20)

When I print Pos, I got the following numbers:

array([1, 3, 2, 3, 1, 2, 1, 2, 2, 3, 6, 0, 0, 4, 0, 1, 1, 3, 1, 5])

You can see from the array that some of the number,such as 6, is generated.

What I want to do it to limit the biggest number(let's say 5), i.e. any random number generated using sct.poisson.rvs should be equal or less than 5,

How can I tweak my code to achieve it. By the way, I am using this in Pandas Dataframe.

Answer 1

What you want could be called the truncated Poisson distribution, except that in the common usage of this term, truncation happens from below instead of from above (example). The easiest, even if not always the most efficient, way to sample a truncated distribution is to double the requested array size and keep only the elements that fall in the desired range; if there are not enough, double the size again, etc. As shown below:

import scipy.stats as sct

def truncated_Poisson(mu, max_value, size):
    temp_size = size
    while True:
        temp_size *= 2
        temp = sct.poisson.rvs(mu, size=temp_size)
        truncated = temp[temp <= max_value]
        if len(truncated) >= size:
            return truncated[:size]

mu = 2.5
max_value = 5
print(truncated_Poisson(mu, max_value, 20))

Typical output: [0 1 4 5 0 2 3 2 2 2 5 2 3 3 3 3 4 1 0 3].