Plot Gamma-Poisson model and Dynamic Gamma-Poisson model
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') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment