Plot random number histogram and PDF(PMF) according the distribution.

plot.py

import argparse
import math
import numpy as np
import matplotlib.pyplot as plt

def plot_exp_pdf(ax, lambda_):
    def exp(lam, x):
        if (x >= 0):
            return lam * np.exp(-lam * x)
        return 0

    xs = np.linspace(0, 20, 2_000)
    ys = [exp(lambda_, x) for x in xs]
    ax.plot(xs, ys, color='C1', label=r'Exponential PDF $\lambda={}$'.format(lambda_))
    ax.set_ylim(0.0)

def plot_poisson_pmf(ax, lambda_):
    def poisson(lam, k):
        return pow(lam, k) * np.exp(-lam) / math.gamma(k+1)

    xs = range(21)
    ys = [poisson(lambda_, x) for x in xs]
    ax.stem(xs, ys, label=r'Poisson PMF $\lambda={}$'.format(lambda_), markerfmt='C1o', linefmt='C1-')
    ax.set_ylim(0.0)

def plot_erlang_kl_pdf(ax, lambda_, k):
    def erlang(lam, k, x):
        return pow(lam, k) * pow(x, k-1) * math.exp(-lam * x) / math.gamma(k)

    xs = np.linspace(0, 20, 2_000)
    ys = [erlang(lambda_, k, x) for x in xs]
    ax.plot(xs, ys, color='C1', label=r'Erlang($E[x]=k/\lambda$) PDF $k={}$, $\lambda={}$'.format(k, lambda_))
    ax.set_ylim(0.0)

def plot_erlang_1l_pdf(ax, lambda_, k):
    def erlang(lam, k, x):
        return pow(k*lam, k) * pow(x, k-1) * math.exp(-lam * k * x) / math.gamma(k)

    xs = np.linspace(0, 20, 2_000)
    ys = [erlang(lambda_, k, x) for x in xs]
    ax.plot(xs, ys, color='C1', label=r'Erlang($E[x]=1/\lambda$) PDF $k={}$, $\lambda={}$'.format(k, lambda_))
    ax.set_ylim(0.0)

parser = argparse.ArgumentParser()
parser.add_argument('--dist')
parser.add_argument('-l', '--lambda_', type=float, default=1.0)
parser.add_argument('-k', '--shape',   type=int,   default=1)
parser.add_argument('-e', '--expectation',         default='')
args = parser.parse_args()

_, ax1 = plt.subplots()
ax2 = ax1.twinx()

xs = np.loadtxt('data.txt')
ax1.hist(xs, align="left", label='Count', bins=20, range=(0, 20))

if args.dist == 'exponential':
    plot_exp_pdf(ax2, args.lambda_)
elif args.dist == 'poisson':
    plot_poisson_pmf(ax2, args.lambda_)
elif args.dist == 'erlang':
    if args.expectation == 'k/l':   # E[x] = k/lambda
        plot_erlang_kl_pdf(ax2, args.lambda_, args.shape)
    elif args.expectation == '1/l': # E[x] = 1/lambda
        plot_erlang_1l_pdf(ax2, args.lambda_, args.shape)

plt.title('{}'.format(args.dist))
plt.legend()
plt.savefig('hist_and_p{}f_{}.png'.format('m' if args.dist == 'poisson' else 'd', args.dist + args.expectation.replace('/', '')))