import numpy as np
import scipy.stats as st
import scipy.optimize as opt
import matplotlib.pyplot as plt

plt.xkcd()
fig = plt.figure( figsize=(6,3) )

# the data:
n = 100
shape = 2.0
scale = 1.0
rng = np.random.RandomState(4637281)
xd = rng.gamma(shape, scale, n)

# true pdf:
xx = np.arange(0.0, 10.1, 0.01)
rv = st.gamma(shape)
yy = rv.pdf(xx)

# mle fit:
a = st.gamma.fit(xd, 5.0)
yf = st.gamma.pdf(xx, *a)

# plot it:
ax = fig.add_subplot( 1, 2, 1 )
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
ax.set_xlim(0, 10.0)
ax.set_ylim(0.0, 0.42)
ax.set_xticks( np.arange(0, 11, 2))
ax.set_yticks( np.arange(0, 0.42, 0.1))
ax.set_xlabel('x')
ax.set_ylabel('Probability density')
ax.plot(xx, yy, '-', lw=5, color='#ff0000', label='pdf')
ax.plot(xx, yf, '-', lw=2, color='#ffcc00', label='mle')
ax.scatter(xd, 0.025*rng.rand(len(xd))+0.05, s=30, zorder=10)
ax.legend(loc='upper right', frameon=False)

# histogram:
h,b = np.histogram(xd, np.arange(0, 8.5, 1), density=True)

# fit histogram:
def gammapdf(x, n, l, s) :
    return st.gamma.pdf(x, n, l, s)
popt, pcov = opt.curve_fit(gammapdf, b[:-1]+0.5*(b[1]-b[0]), h)
yc = st.gamma.pdf(xx, *popt)

# plot it:
ax = fig.add_subplot( 1, 2, 2 )
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
ax.set_xlim(0, 10.0)
ax.set_xticks( np.arange(0, 11, 2))
ax.set_xlabel('x')
ax.set_ylim(0.0, 0.42)
ax.set_yticks( np.arange(0, 0.42, 0.1))
ax.set_ylabel('Probability density')
ax.plot(xx, yy, '-', lw=5, color='#ff0000', label='pdf')
ax.plot(xx, yc, '-', lw=2, color='#ffcc00', label='fit')
ax.bar(b[:-1], h, np.diff(b))
ax.legend(loc='upper right', frameon=False)

plt.tight_layout();
plt.savefig('mlepdf.pdf')
#plt.show();