# EM ALGORITHM FOR SIMPLE 1D GAUSSIAN MIXTURE MODEL.
#
# The arguments are the data vector, the number of components to use, the
# number of EM iterations to do, and a matrix of initial responsibilities 
# of components for data items (default is a random assignment).
#
# The value returned is a list with elements pi (probabilities of components),
# mu (means of components), sigma (standard deviations of components), and
# r, the matrix of responsibilities (which could be passed to another call
# of mix.em).

mix.em <- function (x, K, iters,
                    r = matrix(runif(K*length(x),0.1,1),length(x),K))
{ 
  N <- length(x)

  if (nrow(r)!=N || ncol(r)!=K)
  { stop("Matrix of responsibilities is the wrong size")
  }

  # Make sure initial responsibilities sum to one.

  for (i in 1:N)
  { r[i,] <- r[i,] / sum(r[i,])
  }

  # Do EM iterations, starting with M step, then E step.

  mu <- rep(NA,K)
  sigma <- rep(NA,K)

  for (t in 1:iters)
  {
    cat("\nr:\n")
    print(round(r,3))

    rt <- colSums(r)
   
    # M step.

    pi <- rt / sum(rt)
    for (k in 1:K) 
    { mu[k] <- sum(r[,k]*x) / rt[k]
      sigma[k] <- sqrt (sum(r[,k]*(x-mu[k])^2) / rt[k])
    }

    cat("\npi:",round(pi,3),"  mu:",round(mu,3),"  sigma:",round(sigma,3),"\n")

    # E step.

    for (i in 1:N)
    { r[i,] <- pi * dnorm(x[i],mu,sigma)
      r[i,] <- r[i,] / sum(r[i,]) 
    }
  }

  list (pi=pi, mu=mu, sigma=sigma, r=r)
}