EM算法例子

参考：统计学习方法 EM算法的一个例子_陈嘟嘟cc的博客-CSDN博客_em算法应用实例

# EM algorithm
# coin A, B, C;
# coin A front probability is phi, back probability is 1-phi
# if coin A is front, then choose coin B, coin B front prob is pi, back prob is 1 - pi
# if coin A is back, then choose coin C, coin C front prob is qi, back prob is 1 - qi
# repeat n times (n=10), observation result is [1,1,0,1,0,0,1,0,1,1]
# please estimate param phi, pi, qi

import numpy as np


def em(theta, data, iteration):
    phi = theta[0]
    pi = theta[1]
    qi = theta[2]
    y = np.array(data)
    n = len(data)
    for i in range(iteration):
        # E step: calculate the probability that observation y is from coins B
        mu = phi*pow(pi, y)*pow(1-pi, 1-y) / (phi*pow(pi, y)*pow(1-pi, 1-y) + (1-phi)*pow(qi, y)*pow(1-qi, 1-y))
        # M step: calculate new estimate
        phi = sum(mu) / n
        pi = sum(mu*y) / sum(mu)
        qi = sum((1-mu)*y) / sum(1-mu)
    theta[0] = phi
    theta[1] = pi
    theta[2] = qi
    return theta


if __name__ == "__main__":
    x = np.array([1, 1, 0, 1, 0, 0, 1, 0, 1, 1])
    the = np.array([0.5, 0.5, 0.5])
    res = em(the, x, 10)
    print(res)

运行结果(与初始值设置有关系)：