Search code examples
pythonnumpyfor-loopprobabilitydistribution

Generate two kinds of random variables for each step with same probabilities


goal: Make two kinds of poisson random variables(λ=30, 60)

condition:

  1. We have an array made with the following code(np.random.binomial(n=1, p=1/2, size=1000))
  2. If the number in the array is 1, then we make a poisson random variables(λ=30) else we make a poisson random variables(λ=60)

So I wrote the code like this

import numpy as np
a = np.random.binomial(n=1, p=1/2, size=1000)
b = np.where(a==1, np.random.poisson(lam=30), np.random.poisson(lam=60))
print(b)

But the result was like this

[62 62 28 ... 28 28 62]

Fixed random variables were continually being created.

So how can I get the result that random variables are not fixed? (I don't wanna use loop(for or while))


Solution

  • @BlackRaven's answer works fine, but the use of np.vectorize uses a Python loop "behind the scenes", so it is not as efficient as it could be. An alternative is to use a to create a vector (i.e. 1-d NumPy array) of lambda values to pass to poisson(), like this (note that I've switched to the newer random interface of NumPy):

    rng = np.random.default_rng()
    a = rng.binomial(n=1, p=1/2, size=1000)  # array containing 0 and 1
    lam = 30 + (1-a)*30                      # array containing 30 and 60
    b = rng.poisson(lam=lam)                 # mixture of Poisson samples
    

    Of course, you don't have to create a named variable to hold the lambda values:

    rng = np.random.default_rng()
    a = rng.binomial(n=1, p=1/2, size=1000)  # array containing 0 and 1
    b = rng.poisson(lam=30 + (1-a)*30)       # mixture of Poisson samples
    

    All the loops involved are now performed within NumPy functions that are implemented in C.