Plot Gamma-Poisson model and Dynamic Gamma-Poisson model

dgp.py

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

def pdf_gamma(x, k, lambda_):
    return (math.pow(lambda_, k)/math.gamma(k)) * math.pow(x, k-1) * math.exp(-lambda_*x)

def plot_gamma(ax, alpha, gamma, title):
    plot_x = np.arange(0.01, 1.0, 0.01)
    y = [pdf_gamma(x, alpha, gamma) for x in plot_x ]
    ax.set_title(title)
    ax.axvline(alpha/gamma, color='C1', linestyle='--')
    ax.set_ylim(0, 15)
    ax.plot(plot_x, y, color='C0')

def update(i):
    global alpha, gamma, d_alpha, d_gamma, ctr
    if i != 0:
        axes[0].clear()
        axes[1].clear()
    if i == 20:
        ctr /= 2

    title = r'{}: [{:0>2}] $\lambda$={:.2f}, E(X)={:.2f}'

    X = [np.random.poisson(lam=ctr) for _ in range(args.view) ]
    c = sum(X)
    alpha += c
    gamma += args.view
    plot_gamma(axes[0], alpha, gamma, title.format('Gamma-Poisson', i, ctr, alpha/gamma))

    d_alpha *= args.delta
    d_alpha += c
    d_gamma *= args.delta
    d_gamma += args.view
    plot_gamma(axes[1], d_alpha, d_gamma, title.format(r'Dynamic Gamma-Poisson($\delta$={:.2f})'.format(args.delta), i, ctr, d_alpha/d_gamma))

parser = argparse.ArgumentParser()
parser.add_argument("--lambda_", type=float, default=0.4)
parser.add_argument("--view",    type=int,   default=7)
parser.add_argument("--delta",   type=float, default=0.9)
args = parser.parse_args()

ctr = args.lambda_
alpha = 1
gamma = 1
d_alpha = 1
d_gamma = 1

fig, axes = plt.subplots(nrows=2, sharex=True)
ani = animation.FuncAnimation(fig, update, frames=np.arange(0, 50), interval=200)
ani.save('animate.gif', writer='pillow')