#
# Demonstrate EM estimation of haplotype frequencies
#
mourant <- matrix( c( 91, 147, 85, 
                      32, 78, 75,
                      5, 17, 7), nr=3, byrow=T)
rownames(mourant) <- c("M/M","M/N","N/N")
colnames(mourant) <- c("S/S","S/s","s/s")


oneiter <- function(genotypes, x) {
  
  N <- 2*sum(genotypes)
  oddsr  <-  x[1] * x[4] / (x[3] * x[2])
  LDprop <- oddsr/(1+oddsr)
#
# Expectation step
#
  newtable <- matrix(nr=3, nc=3)
  newtable[1,1] <- x[1] * x[1] / N 
  newtable[1,2] <- 2 * x[1] * x[3] / N
  newtable[1,3] <- x[3] * x[3] / N
  newtable[2,1] <- 2 * x[1] * x[2] / N
  newtable[2,2] <- 2 * (x[1] * x[4] + x[2] * x[3]) / N
  newtable[2,3] <- 2 * x[3] * x[4] / N
  newtable[3,1] <- x[2] * x[2]  / N
  newtable[3,2] <- 2 * x[2] * x[4] / N
  newtable[3,3] <- x[4] * x[4] / N
#
# Maximization step
#
  x[1] <- 2*genotypes[1,1] + genotypes[1,2] + 
        genotypes[2,1] + LDprop*genotypes[2,2]
  
  x[2] <- 2*genotypes[1,3] + genotypes[1,2] + 
        genotypes[2,3] + (1-LDprop)*genotypes[2,2]
  
  x[3] <- 2*genotypes[3,1] + genotypes[3,2] + 
        genotypes[2,1] + (1-LDprop)*genotypes[2,2]
  
  x[4] <- 2*genotypes[3,3] + genotypes[3,2] + 
        genotypes[2,3] + LDprop*genotypes[2,2]
  print(newtable)
# Result
  c(x[1], x[2], x[3], x[4]) 
}
em <- function(g, it=10) {
 ng <- 2*sum(g)
 haplotypes <- c(0.25*ng, 0.25*ng, 0.25*ng, 0.25*ng)
 for(i in 1:it) {
   cat("Iteration ",sprintf("%3d",i), round(haplotypes/ng,3),"\n")
   haplotypes <- oneiter(g, haplotypes)
 }
}