### R commands for Section 4.2
### G. Givens & J. Hoeting, Computational Statistics, 2nd edition, Wiley, 2013.

# BASIC EM ALGORITHM
# Example 4.2 on page 99 in Givens and Hoeting, 2nd edition
#
# Key idea:  use the EM algorithm to estimate the parameters in the
# peppered moth example described on page 99.

#The observed data (phenotype counts for each variety)
nc=85  #Number of observed carbonaria phenotype
ni=196 #number of observed insularia phenotype
nt=341 #Number of observed typica phenotype
n=nc+ni+nt #total sample size

#Initialize the parameter estimate vectors
#One vector for each phenotype
pct=rep(0,20)
pit=rep(0,20)
ptt=rep(0,20)
pct[1]=1/3  #Initial estimates of the parameters
pit[1]=1/3
ptt[1]=1/3

#This for loop carries out EM for the peppered moth example.
for (t in 2:20) {
  #Equations (4.5)-(4.9)
  denom1=(pct[t-1]^2+2*pct[t-1]*pit[t-1]+2*pct[t-1]*ptt[t-1])
  ncct=nc*pct[t-1]^2/denom1
  ncit=2*nc*pct[t-1]*pit[t-1]/denom1
  nctt=nc-ncct-ncit
  niit=ni*pit[t-1]^2/(pit[t-1]^2+2*pit[t-1]*ptt[t-1])
  nitt=ni-niit
  nttt=nt

  #Equations (4.13-4.15)
  pct[t]=(2*ncct+ncit+nctt)/(2*n)
  pit[t]=(2*niit+ncit+nitt)/(2*n)
  ptt[t]=(2*nttt+nctt+nitt)/(2*n)
}

#Examine the output
cbind(pct,pit,ptt)

#equation (4.16) on page 101
rcc=sqrt( (diff(pct)^2+diff(pit)^2)/(pct[-20]^2+pit[-20]^2) )
rcc=c(0,rcc)  #adjusts the length to make the table below

#convergence diagnostics defined below (4.16) on page 101
d1=(pct[-1]-pct[20])/(pct[-20]-pct[20])
d1=c(d1,0)
d2=(pit[-1]-pit[20])/(pit[-20]-pit[20])
d2=c(d2,0)

#Table like Table 4.1 on page 102
print(cbind(pct,pit,rcc,d1,d2)[1:9,],digits=5)
#           pct     pit        rcc       d1      d2
# [1,] 0.333333 0.33333 0.0000e+00 0.042502 0.33659
# [2,] 0.081994 0.23741 5.7069e-01 0.036933 0.18765
# [3,] 0.071249 0.19787 1.6312e-01 0.036727 0.17780
# [4,] 0.070852 0.19036 3.5756e-02 0.036719 0.17624
# [5,] 0.070837 0.18902 6.5860e-03 0.036719 0.17595
# [6,] 0.070837 0.18879 1.1683e-03 0.036719 0.17589
# [7,] 0.070837 0.18875 2.0580e-04 0.036719 0.17588
# [8,] 0.070837 0.18874 3.6205e-05 0.036719 0.17587
# [9,] 0.070837 0.18874 6.3679e-06 0.036736 0.17587

round(c(pct[9], pit[9], 1-pct[9]-pit[9]),3)
# [1] 0.071 0.189 0.740