E step

\[ P(X,\mathbf{z}|\rho, \pi) = \prod_{i=1}^{n} \prod_{p=0}^{P}\left[\pi_{p}N(x_{i}; 0, \Omega_{p})\right]^{\mathbb{I}(z_{i}=p)} \prod_{p=0}^{P}\pi_{p}^{\lambda_{p}-1} \]

\[ \mathbb{E}_{\mathbf{z}|X} \log P(X,\mathbf{z}|\rho, \pi) = \sum_{i=1}^{n} \sum_{p=0}^{P} P(z_{i}=p|X)\left[ \log \pi_{p} + \log N(x_{i}; 0, \Omega_{p})\right] + \sum_{p=0}^{P} (\lambda_{p}-1)\log \pi_{p} \]

\[ \gamma_{z_{i}}(p) = P(z_{i}=p|X_{i}) = \frac{\pi_{p}N(x_{i}; 0, \Omega_{p})}{\sum_{p'=0}^{P}\pi_{p'}N(x_{i}; 0, \Omega_{p'})} \]

M step

\(\pi\): \[ \sum_{i=1}^{n} \gamma_{z_{i}}(p) \frac{1}{\pi_{p}} + \frac{\lambda_{p}-1}{\pi_{p}} - \lambda = 0 \quad \rightarrow \pi_{p} = \frac{1}{\lambda} \left(\sum_{i=1}^{n} \gamma_{z_{i}}(p) + \lambda_{p}-1\right) \quad \lambda = n + \sum_{p=1}^{P}\lambda_{p} - P \]

\[ \hat{\pi}_{p} = \frac{\sum_{i=1}^{n} \gamma_{z_{i}}(p) + \lambda_{p} - 1 }{n + \sum_{p=1}^{P}\lambda_{p} - P } \]

\(\rho\): \[ \begin{align*} f(\rho) &= \sum_{i=1}^{n} \sum_{p=1}^{P} \gamma_{z_{i}}(p)\left[ -\frac{1}{2}\log (1-\phi_{p}^2)-\frac{1}{2(1-\phi_{p}^2)}\left[ \frac{x_{i}^2}{\omega_{p11}^2} + \frac{y_{i}^2}{\omega_{p22}^2} - \frac{2\phi_{p}x_{i}y_{i}}{\omega_{p11}\omega_{p22}}\right] \right]\\ f(\rho)' &= \sum_{i=1}^{n} \sum_{p=1}^{P} \gamma_{z_{i}}(p)\left[ \frac{\phi_{p}}{1-\phi_{p}^2}-\frac{\phi_{p}}{(1-\phi_{p}^2)^2}\left[ \frac{x_{i}^2}{\omega_{p11}^2} + \frac{y_{i}^2}{\omega_{p22}^2}\right] - \frac{\phi_{p}+1}{(1-\phi_{p}^2)^2}\frac{x_{i}y_{i}}{\omega_{p11}\omega_{p22}}\right]\frac{1}{\omega_{p11}\omega_{p22}} = 0 \end{align*} \] \(\phi_p = \frac{\rho + \sigma_{p12}}{\omega_{p11}\omega_{p22}}\), \(\phi_{p}\) is a function of \(\rho\).

Algorithm:

Input: X, Ulist, init_rho, init_pi
Compute loglikelihood
delta = 1
while delta > tol
  E step: update z
  M step: update pi, update rho
  Compute loglikelihood
  Update delta

#' @param rho the off diagonal element of V, 2 by 2 correlation matrix
#' @param Ulist a list of covariance matrices, U_{k}
get_sigma <- function(rho, Ulist){
  V <- matrix(c(1,rho,rho,1), 2,2)
  lapply(Ulist, function(U) U + V)
}

penalty <- function(prior, pi_s){
  subset <- (prior != 1.0)
  sum((prior-1)[subset]*log(pi_s[subset]))
}

#' @title compute log likelihood
#' @param L log likelihoods,
#' where the (i,k)th entry is the log probability of observation i
#' given it came from component k of g
#' @param p the vector of mixture proportions
#' @param prior the weight for the penalty
compute.log.lik <- function(lL, p, prior){
  p = normalize(pmax(0,p))
  temp = log(exp(lL$loglik_matrix) %*% p)+lL$lfactors
  return(sum(temp) + penalty(prior, p))
  # return(sum(temp))
}

normalize <- function(x){
  x/sum(x)
}

mixture.EM.times <- function(X, Ulist, init_rho=0, init_pi=NULL, prior = c('nullbiased', 'uniform'), control = list()){
  times = length(init_rho)
  result = list()
  loglik = c()
  rho = c()
  time.t = c()
  converge.status = c()
  for(i in 1:times){
    out.time = system.time(result[[i]] <- mixture.EM(X, Ulist,
                                                     init_pi=init_pi,
                                                     init_rho=init_rho[i],
                                                     prior=prior,
                                                     control = control))
    time.t = c(time.t, out.time['elapsed'])
    loglik = c(loglik, result[[i]]$loglik)
    rho = c(rho, result[[i]]$rhohat)
    converge.status = c(converge.status, result[[i]]$converged)
  }
  if(abs(max(loglik) - min(loglik)) < 1e-4){
    status = 'global'
  }else{
    status = 'local'
  }
  ind = which.max(loglik)
  return(list(result = result[[ind]], status = status, loglik = loglik, rho=rho, time = time.t, converge.status = converge.status))
}

mixture.EM <- function(X, Ulist, init_rho=0, init_pi = NULL, prior = c('nullbiased', 'uniform'), control = list()) {
  prior = match.arg(prior)
  prior <- mashr:::set_prior(length(Ulist), prior)
  k = length(Ulist)
  if (is.null(init_pi)){
    init_pi <- rep(1/k,k)
  }
  control = ashr:::set_control_squarem(control,nrow(X))
  res = SQUAREM::squarem(par=c(init_pi, init_rho),fixptfn=fixpoint_EM, objfn=negpenloglik,X=X, Ulist=Ulist, prior=prior, control=control)

  return(list(pihat = normalize(pmax(0,head(res$par, -1))), rhohat = tail(res$par, 1), loglik=-res$value.objfn, niter = res$iter, converged=res$convergence, control=control))
}

fixpoint_EM = function(par, X, Ulist, prior){
  rho = tail(par,1)
  pi_s = head(par, -1)
  pi_s = normalize(pmax(0,pi_s)) #avoid occasional problems with negative pis due to rounding

  # compute L
  Sigma <- get_sigma(rho, Ulist)
  L <- t(plyr::laply(Sigma,function(U){mvtnorm::dmvnorm(x=X,sigma=U)}))

  # E
  m  = t(pi_s * t(L)) # matrix_lik is n by k; so this is also n by k
  m.rowsum = rowSums(m)
  classprob = m/m.rowsum #an n by k matrix

  # M
  pinew = normalize(colSums(classprob) + prior - 1)

  rhonew = optimize(EMloglikelihood, interval = c(-1,1), maximum = TRUE, X = X, Ulist = Ulist, z = classprob)$maximum

  return(c(pinew,rhonew))
}

EMloglikelihood = function(rho, X, Ulist, z){
  Sigma = get_sigma(rho, Ulist)
  L = t(plyr::laply(Sigma,function(U){mvtnorm::dmvnorm(x=X,sigma=U, log=TRUE)}))
  sum(L * z)
}

negpenloglik = function(par, X, Ulist, prior){
  Sigma <- get_sigma(tail(par,1), Ulist)
  lL <- t(plyr::laply(Sigma,function(U){mvtnorm::dmvnorm(x=X,sigma=U, log=TRUE)}))
  lfactors <- apply(lL,1,max)
  matrix_llik <- lL - lfactors
  lL = list(loglik_matrix = matrix_llik,
            lfactors   = lfactors)
  ll <- compute.log.lik(lL, head(par, -1), prior)
  return(-ll)
}

Session information

sessionInfo()

R version 3.5.1 (2018-07-02)
Platform: x86_64-apple-darwin15.6.0 (64-bit)
Running under: macOS High Sierra 10.13.6

Matrix products: default
BLAS: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRlapack.dylib

locale:
[1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

other attached packages:
[1] mashr_0.2-15 ashr_2.2-14 

loaded via a namespace (and not attached):
 [1] Rcpp_0.12.19      knitr_1.20        whisker_0.3-2    
 [4] magrittr_1.5      workflowr_1.1.1   REBayes_1.3      
 [7] MASS_7.3-50       pscl_1.5.2        doParallel_1.0.14
[10] SQUAREM_2017.10-1 lattice_0.20-35   foreach_1.4.4    
[13] plyr_1.8.4        stringr_1.3.1     tools_3.5.1      
[16] parallel_3.5.1    grid_3.5.1        R.oo_1.22.0      
[19] rmeta_3.0         git2r_0.23.0      htmltools_0.3.6  
[22] iterators_1.0.10  assertthat_0.2.0  abind_1.4-5      
[25] yaml_2.2.0        rprojroot_1.3-2   digest_0.6.15    
[28] Matrix_1.2-14     codetools_0.2-15  R.utils_2.6.0    
[31] evaluate_0.11     rmarkdown_1.10    stringi_1.2.4    
[34] compiler_3.5.1    Rmosek_8.0.69     backports_1.1.2  
[37] R.methodsS3_1.7.1 mvtnorm_1.0-8     truncnorm_1.0-8