E step

\[ P(\hat{B},B|\rho, \pi) = \prod_{i=1}^{n} \left[N(\hat{b}_{i}; b_{i}, V)\sum_{p=0}^{P} \pi_{p} N(b_{i}; 0, \Sigma_{p})\right] \]

\[ \begin{align*} \mathbb{E}_{B|\hat{B}} \log P(\hat{B},B|\rho, \pi) &= \sum_{i=1}^{n} \mathbb{E}_{b_{i}|\hat{b}_{i}}\left[ \log N(\hat{b}_{i}; b_{i}, V) + \log \sum_{p=0}^{P} \pi_{p} N(b_{i}; 0, \Sigma_{p}) \right] \\ &= \sum_{i=1}^{n} \mathbb{E}_{b_{i}|\hat{b}_{i}}\log N(\hat{b}_{i}; b_{i}, V) + \sum_{i=1}^{n}\mathbb{E}_{b_{i}|\hat{b}_{i}}\log \sum_{p=0}^{P} \pi_{p} N(b_{i}; 0, \Sigma_{p}) \end{align*} \]

\(V\) depends on the first term only. Let \(\mu_{i} = \mathbb{E}_{b_{i}|\hat{b}_{i}}(b_{i})\) \[ \begin{align*} \log N(\hat{b}_{i}; b_{i}, V) &= -\frac{R}{2}\log 2\pi -\frac{1}{2}\log |V| - \frac{1}{2}(\hat{b}_{i}-b_{i})^{T}V^{-1}(\hat{b}_{i}-b_{i}) \\ &= -\frac{R}{2}\log 2\pi -\frac{1}{2}\log |V| - \frac{1}{2}\hat{b}_{i}^{T}V^{-1}\hat{b}_{i} + \frac{1}{2}b_{i}^{T}V^{-1}\hat{b}_{i} + \frac{1}{2}\hat{b}_{i}^{T}V^{-1}b_{i} -\frac{1}{2}b_{i}^{T}V^{-1}b_{i} \\ \mathbb{E}_{b_{i}|\hat{b}_{i}} \log N(\hat{b}_{i}; b_{i}, V) &= -\frac{R}{2}\log 2\pi -\frac{1}{2}\log |V| - \frac{1}{2}\hat{b}_{i}^{T}V^{-1}\hat{b}_{i} + \frac{1}{2}\mu_{i}^{T}V^{-1}\hat{b}_{i} + \frac{1}{2}\hat{b}_{i}^{T}V^{-1}\mu_{i} -\frac{1}{2}tr(V^{-1}\mathbb{E}_{b_{i}|\hat{b}_{i}}(b_{i}b_{i}^{T})) \end{align*} \]

Maximize with respect to V:

We have constraint on V, the diagonal of V must be 1. Let \(V = DCD\), C is the covariance matrix, D = \(diag(1/sqrt(C_{jj}))\).

\[ f(C) = \sum_{i=1}^{n} -\frac{R}{2}\log 2\pi -\frac{1}{2}\log |C|- \log |D| - \frac{1}{2}\hat{b}_{i}^{T}D^{-1}C^{-1}D^{-1}\hat{b}_{i} + \frac{1}{2}\mu_{i}^{T}D^{-1}C^{-1}D^{-1}\hat{b}_{i} + \frac{1}{2}\hat{b}_{i}^{T}D^{-1}C^{-1}D^{-1}\mu_{i} -\frac{1}{2}tr(D^{-1}C^{-1}D^{-1}\mathbb{E}_{b_{i}|\hat{b}_{i}}(b_{i}b_{i}^{T})) \]

\[ \begin{align*} f(C)' &= \sum_{i=1}^{n} -\frac{1}{2}C^{-1} + \frac{1}{2}C^{-1}D^{-1}\hat{b}_{i}\hat{b}_{i}^{T}D^{-1}C^{-1} - \frac{1}{2} C^{-1}D^{-1}\mu_{i}\hat{b}_{i}^{T}D^{-1}C^{-1} - \frac{1}{2}C^{-1}D^{-1}\hat{b}_{i}\mu_{i}^{T}D^{-1}C^{-1} + \frac{1}{2} C^{-1}D^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i})D^{-1}C^{-1} = 0 \\ 0 &= \sum_{i=1}^{n} -\frac{1}{2}C + \frac{1}{2}D^{-1}\hat{b}_{i}\hat{b}_{i}^{T}D^{-1} - \frac{1}{2}D^{-1}\mu_{i}\hat{b}_{i}^{T}D^{-1} - \frac{1}{2}D^{-1}\hat{b}_{i}\mu_{i}^{T}D^{-1} + \frac{1}{2} D^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i})D^{-1} \\ \hat{C} &= \frac{1}{n} \sum_{i=1}^{n} \left[D^{-1}\hat{b}_{i}\hat{b}_{i}^{T}D^{-1} - D^{-1}\mu_{i}\hat{b}_{i}^{T}D^{-1} - D^{-1}\hat{b}_{i}\mu_{i}^{T}D^{-1} + D^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i})D^{-1} \right] \\ &= \frac{1}{n} \sum_{i=1}^{n} \mathbb{E}\left[ (D^{-1}(\hat{b}_{i} - b_{i}))(D^{-1}(\hat{b}_{i} - b_{i}))^{T} | \hat{b}_{i}\right] \\ &= D^{-1}\frac{1}{n} \sum_{i=1}^{n} \mathbb{E}\left[ (\hat{b}_{i} - b_{i})(\hat{b}_{i} - b_{i})^{T} | \hat{b}_{i}\right]D^{-1} \end{align*} \]

We can update C and V as \[ \hat{C}_{(t+1)} = \hat{D}^{-1}_{(t)}\frac{1}{n} \sum_{i=1}^{n} \mathbb{E}\left[ (\hat{b}_{i} - b_{i})(\hat{b}_{i} - b_{i})^{T} | \hat{b}_{i}\right]\hat{D}^{-1}_{(t)} \\ \hat{D}_{(t+1)} = diag(1/\sqrt{\hat{C}_{(t+1)jj}}) \\ \hat{V}_{(t+1)} = \hat{D}_{(t+1)}\hat{C}_{(t+1)}\hat{D}_{(t+1)} \]

The resulting \(\hat{V}_{(t+1)}\) is equivalent as \[ \hat{C}_{(t+1)} = \frac{1}{n}\sum_{i=1}^{n} \mathbb{E}\left[ (\hat{b}_{i} - b_{i})(\hat{b}_{i} - b_{i})^{T} | \hat{b}_{i}\right] \\ \hat{D}_{(t+1)} = diag(1/\sqrt{\hat{C}_{(t+1)jj}}) \\ \hat{V}_{(t+1)} = \hat{D}_{(t+1)}\hat{C}_{(t+1)}\hat{D}_{(t+1)} \]

It is hard to estimate \(\boldsymbol{\pi}\) from the second term, \(\sum_{i=1}^{n}\mathbb{E}_{b_{i}|\hat{b}_{i}}\log \sum_{p=0}^{P} \pi_{p} N(b_{i}; 0, \Sigma_{p})\).

Given V, we estimate \(\boldsymbol{\pi}\) by max loglikelihood, which is a convex problem

Algorithm:

Input: X, Ulist, init_V
Given V, estimate pi by max loglikelihood (convex problem)
Compute loglikelihood
delta = 1
while delta > tol
  M step: update C
  Convert to V
  Given V, estimate pi by max loglikelihood (convex problem)
  Compute loglikelihood
  Update delta

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.MV.times <- function(X, Ulist, init_V = list(diag(ncol(X))), tol=1e-5, prior = c('nullbiased', 'uniform')){
  times = length(init_V)
  result = list()
  loglik = c()
  V = list()
  time.t = c()
  converge.status = c()
  for(i in 1:times){
    out.time = system.time(result[[i]] <- mixture.MV(X, Ulist,
                                                      init_V=init_V[[i]],
                                                      prior=prior,
                                                      tol = tol))
    time.t = c(time.t, out.time['elapsed'])
    loglik = c(loglik, tail(result[[i]]$loglik, 1))
    V = c(V, list(result[[i]]$V))
  }
  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, V=V, time = time.t))
}

mixture.MV <- function(X, Ulist, init_V=diag(ncol(X)), tol=1e-5, prior = c('nullbiased', 'uniform')){
  prior <- match.arg(prior)

  m.model = fit_mash_V(X, Ulist, V = init_V, prior=prior)
  pi_s = get_estimated_pi(m.model, dimension = 'all')
  prior.v <- mashr:::set_prior(length(pi_s), prior)

  # compute loglikelihood
  log_liks <- c()
  log_liks <- c(log_liks, get_loglik(m.model)+penalty(prior.v, pi_s))
  delta.ll <- 1
  niter <- 0
  V = init_V

  while(delta.ll > tol){
    # max_V
    V = E_V(X, m.model)
    V = cov2cor(V)
    m.model = fit_mash_V(X, Ulist, V, prior=prior)
    pi_s = get_estimated_pi(m.model, dimension = 'all')

    log_liks <- c(log_liks, get_loglik(m.model)+penalty(prior.v, pi_s))
    # Update delta
    delta.ll <- log_liks[length(log_liks)] - log_liks[length(log_liks)-1]
    niter <- niter + 1
  }
  return(list(pi = normalize(pi_s), V=V, loglik = log_liks))
}

E_V = function(X, m.model){
  n = nrow(X)
  post.m = m.model$result$PosteriorMean
  post.sec = plyr::laply(1:n, function(i) m.model$result$PosteriorCov[,,i] + tcrossprod(post.m[i,])) # nx2x2 array

  temp1 = crossprod(X)
  temp2 = crossprod(post.m, X) + crossprod(X, post.m)
  temp3 = unname(plyr::aaply(post.sec, c(2,3), sum))

  (temp1 - temp2 + temp3)/n
}

fit_mash_V <- function(X, Ulist, V, prior=c('nullbiased', 'uniform')){
  m.data = mashr::mash_set_data(Bhat=X, Shat=1, V = V)
  m.model = mashr::mash(m.data, Ulist, prior=prior, verbose = FALSE, outputlevel = 3)
  return(m.model)
}

set.seed(1)
n = 4000; p = 2
Sigma = matrix(c(1,0.5,0.5,1),p,p)
U0 = matrix(0,2,2)
U1 = U0; U1[1,1] = 1
U2 = U0; U2[2,2] = 1
U3 = matrix(1,2,2)
Utrue = list(U0=U0, U1=U1, U2=U2, U3=U3)
data = generate_data(n, p, Sigma, Utrue)

m.data = mash_set_data(data$Bhat, data$Shat)
U.c = cov_canonical(m.data)

result.mV <- mixture.MV.times(m.data$Bhat, U.c)

result.mV$result$V

          [,1]      [,2]
[1,] 1.0000000 0.5087773
[2,] 0.5087773 1.0000000

m.data.mV = mash_set_data(data$Bhat, data$Shat, V = result.mV$result$V)
U.c.mV = cov_canonical(m.data.mV)
m.mV = mash(m.data.mV, U.c, verbose= FALSE)
null.ind = which(apply(data$B,1,sum) == 0)

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