Central Limit Theorem Demo

#!/usr/bin/python3
#
#Central Limit Theorem Demo:
# Suppose X_tot = X_1 + ... + X_n, with each X_i following some distribution X_i ~ D,
# then X_tot tends to a normal distribution: X ~N(\mu, \sigma), for n-> \infty
#
#larsupilami73

import os
from scipy import *
from scipy.stats import norm
import matplotlib.pyplot as plt


random.seed(42)
N = 10000 #number of samples to construct


#generate samples from a weird distribution D
#that does not look like a normal distribution
def some_weird_distribution(Nsamples):
    X= []
    n = 0
    while(n<Nsamples):
        x = random.uniform(0,1) #uniform with a hole, couldn't be further from normal...
        if x<0.2 or x>0.8:
            X.append(x)
            n +=1
    return array(X)


#see more color names: plt.cm.colors.cnames
c1 = 'tomato'
c2 = 'steelblue'
c3 = 'ivory'
c4 = 'black'
c5 = 'green'


#temp directory for the frames
if not os.path.exists('./frames'):
    os.mkdir('./frames')
try:
    os.system('rm ./frames/frame*.png')
except:
    pass

#now generate sums of the weird distribution
#and show that the sum approximates a normal distribution
#regardless of what the weird distribution looks like!
fig, (ax1,ax2) = plt.subplots(1,2,figsize=(8,4),dpi=100)

#get the samples from the weird non-gaussian distribution
x = some_weird_distribution(N)
for k in [1,2,3,4,5,6,7,8,9,10,12,14,16,18,20,24,28,32,36,40]:

    #samples
    ax1.scatter(range(N),x, c=c2, s=40,marker='.',alpha=0.4,edgecolor='white', linewidth=0.2,label='K={}\nN={}'.format(k,N))
    ax1.set_title(r'Samples: $X_{tot} = X_1 + \dots +X_{' + '%d'%k +r'}$',fontsize=10)
    ax1.set_ylabel('sample values')
    ax1.set_xlabel('sample nr.')
    ax1.set_ylim(-0.1,1.1)
    ax1.set_yticks([0,0.2,0.4,0.6,0.8,1.0])
    ax1.set_facecolor(c3)
    ax1.legend(loc='upper right',fontsize=9,markerscale=1.4)

    #histogram
    n, bins, patches = ax2.hist(x, bins=50,color=c1,density=True,edgecolor=c4,alpha=0.8)
    ax2.set_title(r'Histogram of $X_{tot}$',fontsize=10)
    ax2.set_ylabel('bin values')
    ax2.set_xlabel('sample values')
    ax2.set_xlim(-0.1,1.1)
    ax2.set_ylim(0,6.5)
    ax2.set_facecolor(c3)

    #add best fit normal to graph
    (mu, sigma) = norm.fit(x)
    y = norm.pdf(bins, mu, sigma)
    label = r'$\mu=%.3f$'%mu
    label +='\n'
    label +=  r'$\sigma=%.3f$' %sigma
    l = ax2.plot(bins, y, 'r--', linewidth=2.5,label=label,color=c4)
    ax2.legend(loc='upper right',fontsize=9,markerscale=0.5,handlelength=0.75)

    #save a frame
    fig.tight_layout()
    fig.savefig('./frames/frame%05d.png' %k)
    ax1.clear()
    ax2.clear()
    #fig.show()

    #add samples of the same distribution to x, to get new samples
    #and renormalization so each vector has the same weight
    x = (k*x + some_weird_distribution(N))/(k+1)

plt.close('all')
#now assemble all the pictures in a gif (requires imagemagick)
os.system('convert -delay 100 -loop 0 ./frames/*.png clt.gif')