参考:统计学习方法 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)
运行结果(与初始值设置有关系):