EM算法-三硬币模型

三硬币模型

 

python实现

import numpy as np
import math
def em_single(theta,Y):
    PB=[]
    pi=theta[0]
    p=theta[1]
    q=theta[2]
    for y in Y:
        u=(pi*math.pow(p,y)*math.pow(1-p,1-y))/(pi*math.pow(p,y)*math.pow(1-p,1-y)+(1-pi)*math.pow(q,y)*math.pow(1-q,1-y))
        PB.append(u)
    new_pi=sum(PB)/len(Y)
    new_p=sum(PB[i]*Y[i] for i in range(len(Y)))/sum(PB)
    new_q=sum((1-PB[i])*Y[i] for i in range(len(Y)))/sum(1-PB[i] for i in range(len(Y)))
    return [new_pi,new_p,new_q]
def em(theta,Y,tol):
    j=0
    while j<1000:
        new_theta=em_single(theta,Y)
        #change=np.abs(new_theta[0] - theta[0])
        # if change<tol:
        if new_theta==theta:
            break
        else:
            theta = new_theta
            j=j+1
        print(new_theta)
    return new_theta
if __name__=="__main__":
    Y = [1, 1, 0, 1, 0, 0, 1, 0, 1, 1]
    theta=[0.4,0.6,0.7]
    result=em(theta,Y,1e-10)
    print("模型参数的极大似然估计是",result)

结果如下图所示：