Python实现EM算法实例代码
脚本专栏 发布日期：2025/11/5 浏览次数：1
正在浏览：Python实现EM算法实例代码
EM算法实例
通过实例可以快速了解EM算法的基本思想，具体推导请点文末链接。图a是让我们预热的，图b是EM算法的实例。
这是一个抛硬币的例子，H表示正面向上，T表示反面向上，参数θ表示正面朝上的概率。硬币有两个，A和B，硬币是有偏的。本次实验总共做了5组，每组随机选一个硬币，连续抛10次。如果知道每次抛的是哪个硬币，那么计算参数θ就非常简单了，如
下图所示：
如果不知道每次抛的是哪个硬币呢？那么，我们就需要用EM算法，基本步骤为：
"text-align: center">
计算过程详解：初始值θ_A^{(0)}θA(0)"color: #ff0000">Python实现
#coding=utf-8
from numpy import *
from scipy import stats
import time
start = time.perf_counter()
def em_single(priors,observations):
"""
EM算法的单次迭代
Arguments
------------
priors:[theta_A,theta_B]
observation:[m X n matrix]
Returns
---------------
new_priors:[new_theta_A,new_theta_B]
:param priors:
:param observations:
:return:
"""
counts = {'A': {'H': 0, 'T': 0}, 'B': {'H': 0, 'T': 0}}
theta_A = priors[0]
theta_B = priors[1]
#E step
for observation in observations:
len_observation = len(observation)
num_heads = observation.sum()
num_tails = len_observation-num_heads
#二项分布求解公式
contribution_A = stats.binom.pmf(num_heads,len_observation,theta_A)
contribution_B = stats.binom.pmf(num_heads,len_observation,theta_B)
weight_A = contribution_A / (contribution_A + contribution_B)
weight_B = contribution_B / (contribution_A + contribution_B)
#更新在当前参数下A，B硬币产生的正反面次数
counts['A']['H'] += weight_A * num_heads
counts['A']['T'] += weight_A * num_tails
counts['B']['H'] += weight_B * num_heads
counts['B']['T'] += weight_B * num_tails
# M step
new_theta_A = counts['A']['H'] / (counts['A']['H'] + counts['A']['T'])
new_theta_B = counts['B']['H'] / (counts['B']['H'] + counts['B']['T'])
return [new_theta_A,new_theta_B]
def em(observations,prior,tol = 1e-6,iterations=10000):
"""
EM算法
：param observations :观测数据
：param prior：模型初值
：param tol：迭代结束阈值
：param iterations：最大迭代次数
：return：局部最优的模型参数
"""
iteration = 0;
while iteration < iterations:
new_prior = em_single(prior,observations)
delta_change = abs(prior[0]-new_prior[0])
if delta_change < tol:
break
else:
prior = new_prior
iteration +=1
return [new_prior,iteration]
#硬币投掷结果
observations = array([[1,0,0,0,1,1,0,1,0,1],
[1,1,1,1,0,1,1,1,0,1],
[1,0,1,1,1,1,1,0,1,1],
[1,0,1,0,0,0,1,1,0,0],
[0,1,1,1,0,1,1,1,0,1]])
print (em(observations,[0.6,0.5]))
end = time.perf_counter()
print('Running time: %f seconds'%(end-start))
总结