Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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
You can’t perform that action at this time.