使用numpy实现期望极大(EM)算法

实现用例为李航统计学习方法第155页例9.1

import numpy as np
"""
通过期望极大算法计算三硬币模型
"""

obs_arr = np.array([1, 1, 0, 1, 0, 0, 1, 0, 1, 1])
num = len(obs_arr)
para_init = {"π":0.4, "p":0.6, "q":0.7}


'''E步，计算在模型参数π、p、q下观测数据y来自掷硬币B的概率'''
def e_func():
    u_j_i_b_arr = np.array([])
    for y_index in obs_arr:
        u_j_i_b = para_init["π"] * pow(para_init["p"], y_index) * pow((1 - para_init["p"]), (1 - y_index))
        u_j_i_c = (1 - para_init["π"]) * pow(para_init["q"], y_index) * pow((1 - para_init["q"]), (1 - y_index))
        u_j_i_b = u_j_i_b / (u_j_i_b + u_j_i_c)
        u_j_i_b_arr = np.append(u_j_i_b_arr, u_j_i_b)
    return u_j_i_b_arr

'''M步，计算模型参数的新估计值'''
def m_func():
    iter_num = 0
    while True:
        π = np.cumsum(e_func())[-1] / num
        # print(π)
        p = np.cumsum(e_func() * obs_arr)[-1] / np.cumsum(e_func())[-1]
        # print(p)
        q = np.cumsum((1 - e_func()) * obs_arr)[-1] / np.cumsum(1 - e_func())[-1]
        # print(q)
        if para_init["π"] == π and para_init["p"] == p and para_init["q"] == q:
            break
        else:
            para_init["π"] = π
            para_init["p"] = p
            para_init["q"] = q
        iter_num +=1
    print("硬币A、B、C正面出现的概率(参数估计值)分别为:")
    print(format(para_init["π"], '.4f'), format(para_init["p"], '.4f'), format(para_init["q"], '.4f'))
    print("迭代的次数:", iter_num)
    
if __name__ == "__main__":
    m_func()