###################################################
#
# Simple EM Example 
#
# Professor M. Zhu
# Current version  = Oct 2023
#
###################################################
rm(list=ls(all=TRUE))

myEMfunc <- function(v, y, theta, max.iter=100, eps=1e-10)
{
# finds MLE for the censored Poisson model,
# given data (v,y), 
# using the EM algorithm instead ...
# written by M. Zhu for teaching illustration

  i = 0; convg<-F; newtheta <- theta
  while (i < max.iter && !convg)
  {
    theta=newtheta
    print(theta)
    # e-step
    ex=ifelse(y>0,(theta*v)/(1-exp(-theta*v)),0)
    # m-step 
    newtheta = sum(ex)/sum(v)
    convg <- abs(theta-newtheta)<eps
    i <- i + 1
  }
  return(list(mle=theta))
}

# Here are some data:
v<-c(8,8,8,8,8,8,8,8,8,8,
     4,4,4,4,4,4,4,4,4,4,
     2,2,2,2,2,2,2,2,2,2,
     1,1,1,1,1,1,1,1,1,1)
y<-c(1,1,1,1,1,1,1,1,1,1,
     1,1,1,1,1,1,1,1,0,0,
     1,1,1,1,1,1,1,0,0,0,
     1,1,1,0,0,0,0,0,0,0)

theta.opt<-myEMfunc(v,y,theta=0.1)