MCMC_fns.R

library( MASS )
library( truncnorm )
library( mvtnorm )
library( MCMCpack )

# data simulation module
sim.CDP.shrink.test = function( lscounts, n, p, m, K, a.er, b.er, sigma.value, sigma.weight, a=2 ){
  # simulate from the prior
  
  # INPUT #
  # lscounts: list of total counts, will generate one dataset for each total counts in the list
  # n: number of biological samples
  # p: number of species
  # m: number of factors
  # K: number of blocks in the underlying covariance matrix
  # allow for misspecification test
  # biological samples are independent between blocks
  # a.er, b.er: hyperparameter in the prior of e_r
  # sigma.value, sigma.weight: discretized prior of small sigma
  # a: random weights of species are generated by sigma*Q^a*(Q>0)
  
  # OUTPUT #
  # List with entries: sigma, Q, X.tru, Y.tru, er, data
  # sigma: vector of simulated small sigma's, length = p
  # Q: matrix of normal latent variables, dim = p*n
  # X.tru, Y.tru: latent factor and loadings, dim(X.tru) = m*p, dim(Y.tru) = m*n
  # er: error term
  # data: a list of matrix, length = length( lscounts )
  # each matrix is p*n.
  
  sigma = sample( sigma.value, p, replace = T, sigma.weight )
  X = matrix( rnorm( p*m ), nrow = m )
  #split the m factors into K blocks
  factor.indx.group = split( 1:m, cut( 1:m, breaks = K ) )
  pop.indx.group = split( 1:n, cut( 1:n, breaks = K ) )
  Y = matrix( 0, nrow = m, ncol = n )
  
  for( x in 1:K){
    raw = matrix( rnorm( length(factor.indx.group[[x]])*length(pop.indx.group[[x]]) ), 
           ncol = length(pop.indx.group[[x]]) )
    Y[factor.indx.group[[x]],pop.indx.group[[x]]] = raw
  }
  
  er = 1/rgamma( 1, a.er, b.er)
  Q = t( apply( t(Y)%*%X, 2, function(x) rnorm( length(x), mean = x, sd = sqrt(er) ) ) )
  final.weights = sigma*(Q*(Q>=0))^2
  data = lapply( lscounts, function(counts) apply( final.weights, 2, function(x) rmultinom( 1, counts, prob = x ) ) )
  
  return( list( sigma = sigma, Q = Q, data = data,
                X.tru=X, Y.tru = Y, er = er ) )
}

# MCMC modules
# univariate M-H with normal approx
get.mode = function( data, sigma, T.aug, mu, s ){
  (mu/s^2 + sqrt((mu/s^2)^2+8*data*(2*sigma*T.aug+1/s^2)))/2/(2*sigma*T.aug+1/s^2)
}
get.s = function( Q, data, sigma, T.aug, mu, s ){
  1/sqrt(2*data/Q^2+2*sigma*T.aug+1/s^2)
}
rej.u = function( Q.prev, Q.prop, data, sigma, T.aug, mu, s, mu.prop, s.prop ){
  2*data*log(Q.prop*(Q.prop>0)) - 
    sigma*T.aug*Q.prop^2*(Q.prop>0) + 
    dnorm( Q.prop, mu, s, log = T ) + 
    dnorm( Q.prev, mu.prop, s.prop, log = T ) -
    2*data*log(Q.prev) + sigma*T.aug*Q.prev^2 - 
    dnorm( Q.prev, mu, s, log = T ) - 
    dnorm( Q.prop, mu.prop, s.prop, log = T )
}

# gibbs step for sampling sigma
sigma.gibbs = function( sigma.old, T.aug, Q, sum.species, sigma.value, sigma.prior, sv.log, sp.log, hold = F ){
  #   browser()
  n = length( T.aug )
  p = length( sigma.old )
  if( hold ){
    sigma.old
  }
  else{
    Q.pos = Q*(Q>0)
    tmp.weights = Q.pos^2%*%T.aug
    vapply( 1:p, function(x) {
      log.w = sv.log*sum.species[x] - sigma.value*tmp.weights[x] + sp.log
      w = exp( log.w - max( log.w )  )
      sample( sigma.value, 1, replace = T, prob = w )
    }, 0 )
  }
}

# Gibbs step for sampling augmented variable T
T.aug.gibbs = function( T.aug.old, sigma, Q, sum.sample, hold = F){
  n = length( T.aug.old )
  if( hold ){
    T.aug.old
  }
  else{
    rgamma( n, shape = sum.sample, rate = sigma%*%(Q^2*(Q>0)) )
  }
}

# Gibbs step for sampling Q
Q.gibbs.ind = function( x, T.aug, Sigma.cond ){
  require( MASS )
  require( truncnorm )
  require( mvtnorm )
  require( MCMCpack )

  #decompress data
  reads = x[1]
  mu.cond = x[2]
  sigma = x[3]
  tmp.Q = x[4]
  
  if( reads==0 ){
    mu.new = mu.cond/(1+2*sigma*T.aug*Sigma.cond^2)
    sigma.new = 1/sqrt(1/Sigma.cond^2+2*sigma*T.aug)
    
    #In real data, this p1 can be really small
    p1.log = pnorm( 0, mu.cond, Sigma.cond, log.p = T ) - pnorm( 0, mu.new, sigma.new, lower.tail = F, log.p = T ) +
      dnorm( 0, mu.new, sigma.new, log = T ) - dnorm( 0, mu.cond, Sigma.cond, log = T )
    p1 = exp( p1.log )
    if( p1 == Inf ){
      p.ratio = 1
    }
    else{
      p.ratio = p1/(1+p1)
    }
    
    if( runif(1)<p.ratio ){
      rtruncnorm( 1, b = 0, mean = mu.cond, sd = Sigma.cond )
    }
    else{
      rtruncnorm( 1, a = 0, mean = mu.new, sd = sigma.new )
    }
  }
  
  #metropolis-hastings for non-zero observation
  else{
    Q.mode = get.mode( reads, sigma, T.aug, mu.cond, Sigma.cond )
    Q.sd = get.s( Q.mode, reads, sigma, T.aug, mu.cond, Sigma.cond )
    Q.prop = rnorm( 1, Q.mode, Q.sd )
    
    if( -rexp(1) < rej.u( tmp.Q, Q.prop, reads, sigma, 
                          T.aug, mu.cond, Sigma.cond, Q.mode, Q.sd ) ){
      Q.prop
    }
    else{
      tmp.Q
    }
  }
}

# Gibbs step for sampling Qij with nij = 0
Q.gibbs.vec.0 = function( para.0 ){
  require( truncnorm )
  
  n.0 = nrow( para.0 )
  mu.new.0 = para.0[,1]/(1+2*para.0[,2]*para.0[,3]*para.0[,4]^2)
  sigma.new.0 = 1/sqrt(1/para.0[,4]^2+2*para.0[,2]*para.0[,3])
  p1.log = pnorm( 0, para.0[,1], para.0[,4], log.p = T ) - pnorm( 0, mu.new.0, sigma.new.0, lower.tail = F, log.p = T ) +
    dnorm( 0, mu.new.0, sigma.new.0, log = T ) - dnorm( 0, para.0[,1], para.0[,4], log = T )
  p1 = exp( p1.log )
  p.ratio = p1/(1+p1)
  p.ratio[is.na(p.ratio)] = 1
  
  label = ( runif( n.0 ) < p.ratio )
  #construct upper and lower vectors
  lower = rep( -Inf, n.0 )
  lower[!label] = 0
  upper = rep( Inf, n.0 )
  upper[label] = 0
  #construct mean and sd vectors
  mean = mu.new.0
  mean[label] = para.0[label,1]
  sd = sigma.new.0
  sd[label] = para.0[label,4]
  
  rtruncnorm( n.0, a = lower, b = upper, mean = mean, sd = sd )
}

# Metropolis step for sampling Qij with nij >0 0
Q.gibbs.vec.n0 = function( para.n0, tmp.Q.n0 ){
  Q.mode = get.mode( para.n0[,1], para.n0[,3], para.n0[,4], para.n0[,2], para.n0[,5] )
  Q.sd = get.s( Q.mode, para.n0[,1], para.n0[,3], para.n0[,4], para.n0[,2], para.n0[,5] )
  n.n0 = length( Q.mode )
  Q.prop = rnorm( n.n0, Q.mode, Q.sd )
  
  #metropolis hastings
  rej = rej.u( tmp.Q.n0, Q.prop, 
                 para.n0[,1], para.n0[,3], para.n0[,4], para.n0[,2], para.n0[,5], 
                 Q.mode, Q.sd )
  labels = ( -rexp(n.n0) < rej )
  
  Q.ret = tmp.Q.n0
  Q.ret[labels] = Q.prop[labels]
  Q.ret
}

######################################################

############# MCMC Sampling begins ###################
main.mcmc.shrink = function( data, start, hyper, 
                             sigma.value, sigma.prior,
                             save_path, 
                             save.obj = c("sigma", "Q", "T.aug", "X", "Y", "er", "delta", "phi"), 
                             burnin = 0.2, thin = 5, step = 1000 ){
  # INPUT
  # data: a count matrix with species in rows and biological samples in column
  # start: starting values of model parameters
  # hyper: values of hyper-parameters in the prior
  # sigma.value, sigma.prior: discretized sigma prior
  # save_path: folder to save the MCMC results
  # save.obj: list of model parameters that will be saved, default is all parameters
  # burnin: burn-in fraction
  # thin: thinning size
  # step: number of iterations of MCMC chain
  
  # OUTPUT
  # In the save_path, there will be many R RDS files, each is named res_[iteration_i] and records
  # the posterior samples of save.obj at iteration_i
  
  require( "mvtnorm" )
  nv = hyper$nv
  a.er = hyper$a.er
  b.er = hyper$b.er
  a1 = hyper$a1
  a2 = hyper$a2
  m = hyper$m
  
  n = ncol( data )
  p = nrow( data )
  sum.species = rowSums( data )
  sum.sample = colSums( data )
  sv.log = log( sigma.value )
  sp.log = log( sigma.prior )
  
  sigma.old = as.vector( start$sigma )
  Q.old = start$Q
  T.aug.old = as.vector( start$T.aug )
  X.old = start$X
  Y.old = start$Y
  er.old = start$er
  delta.old = as.vector(start$delta)
  phi.old = start$phi
  
  all.cache = c()
  t.t = proc.time()
  for( iter in 2:step ){
#     browser()
    if( iter %% 1 == 0 )
      print( iter )
    #sample sigma
    #this is a major time consuming step
    all.cache$sigma = sigma.old
    sigma.old = sigma.gibbs( 
      sigma.old, T.aug.old, 
      Q.old, sum.species, 
      sigma.value, sigma.prior,
      sv.log, sp.log
    )

    #sample T
    all.cache$T.aug = T.aug.old
    T.aug.old = T.aug.gibbs( 
      T.aug.old, sigma.old, 
      Q.old, sum.sample
    )
    
    #sample Q
    #matrix of all parameters
    Q.para.all = cbind( c(data), c(t(X.old)%*%Y.old), 
                        rep( sigma.old, n ), c(Q.old),
                        rep( T.aug.old, each = p ), rep( sqrt(er.old), n*p ) )
    all.cache$Q = Q.old
    #get samples
    labels = (Q.para.all[,1]==0)
    Q.0 = Q.gibbs.vec.0( Q.para.all[labels, c(2,3,5,6)] )
    Q.n0 = Q.gibbs.vec.n0( Q.para.all[!labels, c(1,2,3,5,6)], 
                           Q.para.all[!labels, 4] )
    #save it into proper locations
    Q.tmp = rep( 0, nrow( Q.para.all ) )
    Q.tmp[labels] = Q.0
    Q.tmp[!labels] = Q.n0
    #convert it back to matrix
    Q.old = matrix( Q.tmp, nrow = p )
    
    #sample Y
    #this is a major time consuming step
    #get tau
    tau.tmp = cumprod( delta.old )
    all.cache$Y = Y.old
    Y.old = vapply( 1:n, function(x){
      Sigma.Y = solve( X.old%*%t(X.old)/er.old + 
                         diag( phi.old[x,]*tau.tmp ) )
      mu.Y = Sigma.Y%*%X.old%*%Q.old[,x]/er.old
      rmvnorm( 1, mu.Y, sigma = Sigma.Y )
    }, rep( 0, nrow( Y.old ) ) )
    
    #sample er
    #prior is ga(1,10)
    all.cache$er = er.old
    er.old = 1/rgamma( 1, shape = n*p/2+a.er, 
                       rate = sum((Q.old-t(X.old)%*%Y.old)^2)/2+b.er )
    
    #sample X
    Sigma.X = solve( Y.old%*%t(Y.old)/er.old + diag( m ) )
    all.cache$X = X.old
    X.mean = Sigma.X%*%Y.old%*%t(Q.old)/er.old
    X.old = X.mean + t( rmvnorm( p, sigma = Sigma.X ) )
    
    #sample phi
    all.cache$phi = phi.old
    phi.old = matrix( rgamma( n*m, shape = (nv+1)/2, rate = c( Y.old^2*tau.tmp + nv )/2 ), 
                      nrow = n, byrow = T )
    
    #sample delta
    delta.tmp = delta.old
    all.cache$delta = delta.old
    for( k in 1:m ){
      if( k == 1 ){
        delta.prime = cumprod( delta.tmp )/delta.tmp[1]
        delta.tmp[k] = rgamma( 1, shape = n*m/2 + a1, 
                               rate = 1 + sum(colSums( t(Y.old^2)*phi.old )*delta.prime)/2
        )
      }
      else{
        delta.prime = cumprod( delta.tmp )[-(1:(k-1))]/delta.tmp[k]
        delta.tmp[k] = rgamma( 1, shape = n*(m-k+1)/2 + a2, 
                               rate = 1 + sum( colSums( t(Y.old^2)*phi.old )[-(1:(k-1))]*delta.prime)/2
        )
      }
    }
    delta.old = delta.tmp
    
    if( iter > burnin*step & iter%%thin == 0 )
    saveRDS( all.cache[save.obj], file = paste( save_path, iter-1, sep = "_") )
  }
  
  #save last run
  all.cache = list( sigma = sigma.old, T.aug = T.aug.old, Q = Q.old, Y = Y.old, er = er.old,
                    X = X.old, phi = phi.old, delta = delta.old )
  saveRDS( all.cache[save.obj], file = paste( save_path, iter, sep = "_") )
}