# compute log-likelihood for tracking
function log_likelihood(y, p, mu, sigma)
N = length(y)
K = length(p)
ll_comp = zeros(N, K)
for k = 1:K
ll_comp[:, k] = p[k] * pdf.(Normal(mu[k], sigma[k]), y)
end
ll = sum(log.(sum(ll_comp, dims=2)))
return ll
end
function em_algorithm(y, p_init, mu_init, sigma_init; max_iter = 300, thresh=1e-4)
N = length(y)
K = length(p_init)
# initialize storage and add initial guesses
p = zeros(max_iter + 1, K)
p[1, :] = p_init
mu = zeros(max_iter + 1, K)
mu[1, :] = mu_init
sigma = zeros(max_iter + 1, K)
sigma[1, :] = sigma_init
# compute initial log-likelihood
ll = zeros(max_iter + 1)
ll[1] = log_likelihood(y, p[1, :], mu[1, :], sigma[1, :])
iter = 1
for t = 1:max_iter
w = e_step(y, p[t, :], mu[t, :], sigma[t, :])
p[t+1, :], mu[t+1, :], sigma[t+1, :] = m_step(y, w)
ll[t+1] = log_likelihood(y, p[t+1, :], mu[t+1, :], sigma[t+1, :])
if abs(ll[t+1] - ll[t]) < thresh
break
end
iter += 1
end
return (p[1:iter+1, :], mu[1:iter+1, :], sigma[1:iter+1, :], ll[1:iter+1, :])
end