Univariate Poisson Mixture Model

univariate_pmm.py

import numpy as np
import random as rand

from scipy.stats import poisson

SEED = 42
N = 100

def generate():
    import matplotlib.pyplot as plt

    rand.seed(SEED)
    np.random.seed(SEED)

    params = {'lambda1': 1, 'lambda2': 5, 'theta':0.40};
    print(params);
    n1 = int(N * params['theta']);
    n2 = N - n1;
    print("n1 = {}, n2 = {}".format(n1,n2));

    p1 = np.random.poisson(params['lambda1'], n1)
    p2 = np.random.poisson(params['lambda2'], n2)
    data = np.hstack([p1, p2])
    plt.hist(data, bins=15);
    plt.show()

    return data;

def E(data, l1, l2, theta):
    prob = []
    p1 = poisson(l1)
    p2 = poisson(l2)
    for i, d in enumerate(data):
        prob.append([p1.pmf(d)*theta , p2.pmf(d) * (1.0 - theta) ])
    return np.array(prob)


def M(data, E):
    a = 1.0 * (E[:,0] > E[:,1])
    #print("a={}".format(a))
    d1 = []
    d2 = []
    p1 = []
    p2 = []
    for i, d in enumerate(data):
        p = ( E[i, 0] / (E[i, 0] + E[i, 1]))
        d1.append(d * p)
        d2.append(d * ( 1.0 - p))
        p1.append(p)
        p2.append(1.0 - p)
    d1 = np.array(d1);
    d2 = np.array(d2);
    assert len(data) == len(d1) == len(d2)
    theta = np.sum(p1)/len(data)
    l1 = np.sum(d1/sum(p1))
    l2 = np.sum(d2/sum(p2))
    return l1,l2,theta

def initialize_params(data):
    theta = data.mean()/data.max()

    return min(data)+1, max(data), theta;

if '__main__' == __name__:
    data = generate()
    print(data)
    lambda1, lambda2, theta = initialize_params(data)
    print("Initialized params: lambda1={}, lambda2={}, theta={}".format(lambda1, lambda2, theta))
    for i in range(20):
        print("="*20)
        print("i={}".format(i))
        prob = E(data, lambda1, lambda2, theta)
        lambda1, lambda2, theta = M(data, prob)
        print("lambda1={}, lambda2={}, theta={}".format(lambda1, lambda2, theta))