Search code examples
pythonrandomnumpyscipysampling

Vectorised code for sampling from truncated normal distributions with different intervals


The following code generates a sample of size 100 from trunctated normal distributions with different intervals. Is there any effecient(vectorised) way of doing this?

from scipy.stats import truncnorm
import numpy as np
sample=[]
a_s=np.random.uniform(0,1,size=100)
b_s=a_s+0.2
for i in range(100):
    sample.append(truncnorm.rvs(a_s[i], b_s[i], size=100))
print sample

Solution

  • One day in the not so distant future, all NumPy/SciPy functions will broadcast all their arguments, and you will be able to do truncnorm.rvs(a_s, b_s, size=100), but since we are not there yet, you could manually generate your random samples from a uniform distribution and the CDF and PPF of a normal distribution:

    import numpy as np
    from scipy.stats import truncnorm, norm
    
    a_s = np.random.uniform(0, 1, size=100)
    b_s = a_s + 0.2
    
    cdf_start = norm.cdf(a_s)
    cdf_stop = norm.cdf(b_s)
    cdf_samples = np.random.uniform(0, 1, size=(100, 100))
    cdf_samples *= (cdf_stop - cdf_start)[:, None]
    cdf_samples +=  cdf_start[:, None]
    truncnorm_samples = norm.ppf(cdf_samples)