# em07.r
# xref: em05.tex, em07.c
# input:
# output:
# does:  Mixture of 2 normals with equal known variances.
#      Z ~ Bern(theta)
#      X0 ~ N(mu0, 1)
#      X1 ~ N(mu1, 1)
#      mu0 < mu1
#      - MLE using EM algorithm
#
#***************************************************/

main = function (dummy)

{
  truparms = list(mu0 = 0, mu1 = 3, theta = 0.3);
  n        = 100;

  y = simulate (truparms, n);          # /* simulate data */
  estparms = estimate (y)$parms;       # /* estimate parameters  */
  summary  (truparms, estparms, y);
}
#***************************************************/

estimate = function (y)

# input: data y[1:n]
# output: conv, parms
{
  tol     = 1e-9;
  maxiter = 200;

  if (length(y) < 3) return (list(conv = FALSE, parms = 0));

  parms = initvals (y);           # /* initial values */
  pll   = monitor (y, parms, 0);  # /* previous ll */

  for (i in 1:maxiter)  {
    esave = Estep  (y, parms);
    parms = Mstep  (esave);
    ll    = monitor (y, parms, i);
    conv  =  (abs(ll-pll) < tol);   # /* converged ? */
    if (conv) break;
    pll   = ll;
  }

  list(conv = conv, parms = parms)
}

#***************************************************/

Estep = function (y, parms)

#  w  = E[Z  | Y]
#  c0 = E[X0 | Y]
#  c1 = E[X1 | Y]
{
  d = parms$mu1 - parms$mu0;
  s = 0.5*(parms$mu1 + parms$mu0);
  priorodds = parms$theta/(1-parms$theta);

  psi = priorodds * exp(d*(y-s));
  w   = psi/(1+psi);
  c0  = w*parms$mu0 + (1-w)*y;
  c1  = w*y + (1-w)*parms$mu1;

  list(w = w, c0 = c0, c1 = c1)
}
#***************************************************/

Mstep = function (esave)

# obtain new parameter estimates
{

  mu0   = mean(esave$c0);
  mu1   = mean(esave$c1);
  theta = mean(esave$w);

  if (mu0 > mu1) {
    tmp=mu0; mu0=mu1; mu1=tmp;
    theta = 1 - theta;
  }

  list(mu0 = mu0, mu1 = mu1, theta = theta)
}
#***************************************************/

loglkhd = function (y, parms)

# the log-likelihood  based on y[1:n]
{
  f = parms$theta       * dnorm(y-parms$mu1) +
      (1-parms$theta)   * dnorm(y-parms$mu0);

  sum(log(f))
}
#***************************************************/

initvals = function (y)

# parms = initial values based on y[1:n]
{
  ybar = mean(y);
  sd   = sd(y);

  list(mu0 = ybar-sd, mu1 = ybar+sd, theta = 0.5);
}
#***************************************************/

monitor = function (y, parms, iter)

#  /* print log-lkhd and estimates */
{
  ll = loglkhd (y, parms);

  print (paste('Iteration=', iter,
               'l=',     ll,
               'theta=', parms$theta,
               'mu0=',   parms$mu0,
               'mu1=',   parms$mu1 ) );

  ll
}
#***************************************************/

summary = function (truparms, estparms, y)

{
  print (paste('Sample size n=', length(y)));

  print (paste('True parameter values:',
    'theta=', truparms$theta,
    'mu0=',   truparms$mu0,
    'mu1=',   truparms$mu1 ));

  print (paste('Parameter Estimates:',
    'theta=', estparms$theta,
    'mu0=',   estparms$mu0,
    'mu1=',   estparms$mu1 ));

  print (paste('Log-likelihood = ', loglkhd(y, estparms)));

}
#***************************************************/

simulate = function (parms, n)

# input: parms, n (ignored)
# output: y[1:n], column vec, fixed data set.
{
  c ( #  /* simulated with parms: 0, 2, 0.3 */
    0.0007668, -0.939475, 0.8456077, -0.493304, 0.3805769, 2.6011146,
    2.08912, 0.1207658, 2.1957545, -0.959282, 0.1548358, 0.3394777,
    0.4788278, 1.1161578, 0.7300012, -2.120032, -1.085473, 0.2274149,
    0.8690178, -0.932938, 0.4580891, 1.2779732, -1.400311, 3.0479398,
    -1.200237, 2.4490362, 1.1866626, -0.693796, -1.651684, -2.222105,
    0.1401175, 1.9019002, 0.6396009, -1.382346, -0.718541, -0.279965,
    0.4651524, 4.3747447, 1.3001852, -1.359992, 0.2648149, 1.524874,
    1.6180265, 1.7134072, -0.870176, -2.099657, 0.3258149, 1.9970998,
    3.4406849, 0.5807176, -0.93498, 1.2195432, 1.0252208, 2.7270676,
    -1.145333, 0.9615463, -0.595377, 0.320387, 3.0322296, 0.0399386,
    0.6169643, 1.0346161, -1.311096, 0.7927389, -0.37912, 1.2072907,
    0.3468948, 0.5203926, -0.932413, -0.597823, -0.854034, 2.5486596,
    -0.502315, 3.2556189, -0.598408, 0.3365086, 0.4282481, 2.5592742,
    1.9806215, -0.912323, 0.4092242, -0.365883, 0.9473954, 0.4847945,
    0.0288647, 0.4406987, -1.126491, 4.0417427, -1.564552, 2.8244872,
    1.8998997, 3.807505, -0.419282, 2.7611861, 2.0530032, 1.2842716,
    0.8478354, -1.082433, 2.3424769, 0.2944996
  )
}
#***************************************************/

simulate2 = function (parms, n)

#/* input: parms, n */
#/* output: data y[1:n] simulated from the mixture model */
{
  z  = (runif(n) <= parms$theta);
  mu = (1-z)*parms$mu0 + z*parms$mu1;

  mu + rnorm(n)
}
#****************************************************/