Estimate cor — M V decompose

E step

$$p(\hat{\mathbf{B}}, \mathbf{B}, \mathbf{z}) h(\boldsymbol{\pi}) = \prod_{j=1}^{J} \prod_{p = 1}^{P}\left[\pi_{p} N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{b}_{j}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j})N_{R}(\mathbf{b}_{j}; \mathbf{0}, \Sigma_{p})\right]^{\mathbb{I}(z_{j} = p)} \prod_{p=1}^{P} \pi_{p}^{\lambda_{p}-1}$$

$$\begin{align*} P(z_{j}=p, \mathbf{b}_{j}|\hat{\mathbf{b}}_{j}) &= \frac{P(z_{j}=p, \mathbf{b}_{j},\hat{\mathbf{b}}_{j})}{P(\hat{\mathbf{b}}_{j})} = \frac{P(\hat{\mathbf{b}}_{j}|\mathbf{b}_{j})P(\mathbf{b}_{j}|z_{j}=p) P(z_{j}=p)}{P(\hat{\mathbf{b}}_{j})} \\ &= \frac{\pi_{p} N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{b}_{j}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j})N_{R}(\mathbf{b}_{j}; \mathbf{0}, \Sigma_{p})}{\sum_{p'}\pi_{p'} N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{0}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j} + \Sigma_{p'})} \\ &= \frac{\pi_{p} N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{0}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j} + \Sigma_{p})}{\sum_{p'}\pi_{p'} N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{0}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j} + \Sigma_{p'})} \frac{N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{b}_{j}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j})N_{R}(\mathbf{b}_{j}; \mathbf{0}, \Sigma_{p})}{N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{0}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j} + \Sigma_{p})} \\ &= \gamma_{jp} P(\mathbf{b}_{j}|z_{j}=p, \hat{\mathbf{b}}_{j}) \\ &= P(z_{j}=p|\hat{\mathbf{b}}_{j}) P(\mathbf{b}_{j}|z_{j}=p, \hat{\mathbf{b}}_{j}) \end{align*}$$

E step: $$\begin{align*} \mathbb{E}_{\mathbf{z}, \mathbf{B}|\hat{\mathbf{B}}}\log p(\hat{\mathbf{B}}, \mathbf{B}, \mathbf{z}) h(\boldsymbol{\pi}) &= \mathbb{E}_{\mathbf{z}, \mathbf{B}|\hat{\mathbf{B}}} \{ \sum_{j=1}^{J}\sum_{p = 1}^{P} \mathbb{I}(z_{j} = p)\left[\log \pi_{p} + \log N_{R}(\hat{\mathbf{b}}_{j}; \mathbf{b}_{j}, \mathbf{S}_{j}\mathbf{V}\mathbf{S}_{j}) + \log N_{R}(\mathbf{b}_{j}; \mathbf{0}, \Sigma_{p})\right] + \sum_{p=1}^{P} (\lambda_{p}-1) \log \pi_{p} \} \\ &= \sum_{j=1}^{J} \sum_{p=1}^{P} \gamma_{jp} \left[\log \pi_{p} - \frac{1}{2}\log |\mathbf{V}| - \log |\mathbf{S}_{j}| - \frac{1}{2}\mathbb{E}_{\mathbf{b}_{j}|\hat{\mathbf{b}}_{j}, z_{j}=p}\left((\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})^{T}\mathbf{S}_{j}^{-1}\mathbf{V}^{-1}\mathbf{S}_{j}^{-1}(\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})\right) - \frac{1}{2}\log |\Sigma_{p}| - \frac{1}{2}\mathbb{E}_{\mathbf{b}_{j}|\hat{\mathbf{b}}_{j}, z_{j}=p}\left(\mathbf{b}_{j}^{T}\Sigma_{p}^{-1}\mathbf{b}_{j} \right) \right] + \sum_{p=1}^{P} (\lambda_{p}-1)\log \pi_{p} \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}))$.

$$\begin{align*} f(\mathbf{C}) &= \sum_{j=1}^{J} \sum_{p=1}^{P} \gamma_{jp} \left[- \frac{1}{2}\log |\mathbf{D}\mathbf{C}\mathbf{D}| - \frac{1}{2}\mathbb{E}_{\mathbf{b}_{j}|\hat{\mathbf{b}}_{j}, z_{j}=p}\left((\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})^{T}\mathbf{S}_{j}^{-1}\mathbf{D}^{-1}\mathbf{C}^{-1}\mathbf{D}^{-1}\mathbf{S}_{j}^{-1}(\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})\right) \right] \\ &= \sum_{j=1}^{J} \sum_{p=1}^{P} \gamma_{jp} \left[- \frac{1}{2}\log |\mathbf{C}| - \log |\mathbf{D}|- \frac{1}{2}\mathbb{E}_{\mathbf{b}_{j}|\hat{\mathbf{b}}_{j}, z_{j}=p}\left((\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})^{T}\mathbf{S}_{j}^{-1}\mathbf{D}^{-1}\mathbf{C}^{-1}\mathbf{D}^{-1}\mathbf{S}_{j}^{-1}(\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})\right) \right] \\ f(\mathbf{C})' &= \sum_{j=1}^{J} \sum_{p=1}^{P} \gamma_{jp}\left[ -\frac{1}{2} \mathbf{C}^{-1} + \frac{1}{2} \mathbf{C}^{-1} \mathbf{D}^{-1}\mathbf{S}_{j}^{-1}\mathbb{E}\left((\hat{\mathbf{b}}_{j}-\mathbf{b}_{j}) (\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})^{T}|\hat{\mathbf{b}}_{j}, z_{j} = p \right)\mathbf{S}_{j}^{-1}\mathbf{D}^{-1} \mathbf{C}^{-1} \right] = 0 \\ \mathbf{C} &= \frac{1}{J} \sum_{j=1}^{J} \sum_{p=1}^{P} \gamma_{jp}\mathbf{D}^{-1}\mathbf{S}_{j}^{-1}\mathbb{E}\left((\hat{\mathbf{b}}_{j}-\mathbf{b}_{j}) (\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})^{T}|\hat{\mathbf{b}}_{j}, z_{j} = p \right)\mathbf{S}_{j}^{-1}\mathbf{D}^{-1} \\ &= \frac{1}{J} \mathbf{D}^{-1}\sum_{j=1}^{J} \mathbf{S}_{j}^{-1}\mathbb{E}\left((\hat{\mathbf{b}}_{j}-\mathbf{b}_{j}) (\hat{\mathbf{b}}_{j}-\mathbf{b}_{j})^{T}|\hat{\mathbf{b}}_{j}\right)\mathbf{S}_{j}^{-1}\mathbf{D}^{-1} \end{align*}$$

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

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

The estimated V maximize the lower bound of the log likelihood. We estimate $\boldsymbol{\pi}$ from the log likelihood.

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]))
}

mixture.MV <- function(mash.data, Ulist, init_V=diag(ncol(mash.data$Bhat)), max_iter = 500, tol=1e-5, prior = c('nullbiased', 'uniform'), cor = TRUE, track_fit = FALSE){
  prior <- match.arg(prior)
  tracking = list()

  m.model = fit_mash_V(mash.data, 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 <- numeric(max_iter+1)
  log_liks[1] <- get_loglik(m.model)+penalty(prior.v, pi_s)
  V = init_V
  
  result = list(V = V, logliks = log_liks[1], mash.model = m.model)
  
  for(i in 1:max_iter){
    if(track_fit){
      tracking[[i]] = result
    }
    # max_V
    V = E_V(mash.data, m.model)
    if(cor){
        V = cov2cor(V)
    }
    m.model = fit_mash_V(mash.data, Ulist, V, prior=prior)
    pi_s = get_estimated_pi(m.model, dimension = 'all')

    log_liks[i+1] <- get_loglik(m.model)+penalty(prior.v, pi_s)
    
    result = list(V = V, logliks = log_liks[1:(i+1)], mash.model = m.model)

    # Update delta
    delta.ll <- log_liks[i+1] - log_liks[i]
    if(delta.ll<=tol) break;
  }
  
  if(track_fit){
    result$trace = tracking
  }
  
  return(result)
}

E_V = function(mash.data, m.model){
  n = mashr:::n_effects(mash.data)
  Z = mash.data$Bhat/mash.data$Shat
  post.m.shat = m.model$result$PosteriorMean / mash.data$Shat
  post.sec.shat = plyr::laply(1:n, function(i) (t(m.model$result$PosteriorCov[,,i]/mash.data$Shat[i,])/mash.data$Shat[i,]) + tcrossprod(post.m.shat[i,])) # nx2x2 array
  temp1 = crossprod(Z)
  temp2 = crossprod(post.m.shat, Z) + crossprod(Z, post.m.shat)
  temp3 = unname(plyr::aaply(post.sec.shat, c(2,3), sum))

  (temp1 - temp2 + temp3)/n
}

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

Data

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(m.data, U.c)

The estimated $V$ is

result.mV$V

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

m.mV = result.mV$mash.model
null.ind = which(apply(data$B,1,sum) == 0)

The log likelihood is -12302.52. There are 26 significant samples, 0 false positives. The RRMSE is 0.5822283.

Session information

sessionInfo()

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

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.18.0454 ashr_2.2-7       

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.18    
[28] Matrix_1.2-14     codetools_0.2-15  R.utils_2.7.0    
[31] evaluate_0.12     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