R/vbvs_concurrent.R

In jeff-goldsmith/vbvs.concurrent: Variational Bayes Variable Selection for the Functional Linear Concurrent Model

#' vbvs_concurrent
#' 
#' Implements variational bayes variable selection for the linear functional concurrent model
#' 
#' @param formula formula for desired regression. should have form \code{ y ~ x1 + x2 + ... + x_k | t}, where \code{t} is the variable that parameterized observed functions
#' @param id.var variable giving subject ID vector
#' @param data optional data frame
#' @param Kt number of spline basis functions for coefficients and FPCs
#' @param Kp number of FPCs to estimate
#' @param v0 tuning parameter; normal spike variance
#' @param v1 tuning parameter; normal slab variance
#' @param standardized logical; are covariates already standardized?
#' @param t.min minimum value to be evaluated on the time domain (useful if data are sparse and / or irregular). if `NULL`, taken to be minium observed value.
#' @param t.max maximum value to be evaluated on the time domain (useful if data are sparse and / or irregular). if `NULL`, taken to be minium observed value.
#' 
#' @author Jeff Goldsmith \email{jeff.goldsmith@@columbia.edu}
#' 
#' @importFrom stats model.frame
#' @importFrom splines bs
#' 
#' @export
#' 
#' @examples
#' \dontrun{
#' library(dplyr)
#' library(tidyr)
#' library(refund.shiny)
#' 
#' ## set design elements
#' set.seed(1)
#' I = 100
#' p = 50
#' 
#' ## coefficient functions
#' beta1 = function(t) { sin(2*t*pi) }
#' beta2 = function(t) { cos(2*t*pi) }
#' beta3 = function(t) { 1 }
#' 
#' ## FPC basis functions
#' psi1 = function(t) { sin(2*t*pi) }
#' psi2 = function(t) { cos(2*t*pi) }
#' 
#' ## generate subjects, observation times, and FPC scores
#' time.data = sapply(1:I, function(u) {
#'   ji = sample(10:15, 1)
#'   rbind(runif(ji, 0, 1) %>% sort, 
#'         rep(u, ji),
#'         rep(rnorm(1, 0, 3), ji),
#'         rep(rnorm(1, 0, 1), ji))
#' }) %>% unlist() %>% matrix(ncol = 4, byrow = TRUE)
#' colnames(time.data) = c("time", "subj", "c_i1", "c_i2")
#' time.data = as.data.frame(time.data)
#' 
#' ## generate predictor data
#' predictor.data = matrix(rnorm(dim(time.data)[1] * p), dim(time.data)[1], p)
#' colnames(predictor.data) = paste0("Cov_", 1:p)
#' 
#' ## combine and generate responses
#' concurrent.data = cbind(time.data, predictor.data)
#' concurrent.data = 
#'   mutate(concurrent.data,
#'          Y = Cov_1 * beta1(time) +            ## fixed effects
#'              Cov_2 * beta2(time) + 
#'              Cov_3 * beta3(time) + 
#'              c_i1 * psi1(time) +              ## pca effects
#'              c_i2 * psi2(time) + 
#'              rnorm(dim(concurrent.data)[1]))  ## measurement error
#' 
#' ## fit model
#' pred.list = paste("Cov", 1:p, sep = "_")
#' formula = as.formula( paste("Y ~", paste(pred.list, collapse = "+"), "| time") )
#' fit.vb = vb_concurrent(formula, id.var = "subj", data = concurrent.data, standardized = TRUE,
#'                            t.min = 0, t.max = 1, Kp = 3)
#' fit.vbvs = vbvs_concurrent(formula, id.var = "subj", data = concurrent.data, standardized = TRUE,
#'                            t.min = 0, t.max = 1, Kp = 3)
#' 
#' ## interactive plot for mean model
#' plot_shiny(fit.vb)
#' plot_shiny(fit.vbvs)
#'
#' ## interactive plot for fpca expansion of residuals
#' plot_shiny(fit.vbvs$fpca.obj)
#' }
vbvs_concurrent = function(formula, id.var = NULL, data=NULL, Kt = 5, Kp = 2, v0 = 0.01, v1 = 100,
                           standardized = FALSE, t.min = NULL, t.max = NULL){
  
  ## parse formula
  LHS = as.character(formula)[2]
  RHS = as.character(formula)[3]
  
  if (is.null(id.var)) { stop("Please specify the subject ID variable")}
  if (grep("|", RHS) == 0) { stop("Formula incorrectly specified; needs a '|' to indicate time parameter")}
  
  time.var = strsplit(RHS, "|", fixed = TRUE)[[1]][2] %>% gsub(" ", "", x = ., fixed = TRUE)
  RHS = strsplit(RHS, "|", fixed = TRUE)[[1]][1]
  RHS.mod = paste(RHS, "+", time.var, "+", id.var)
  
  formula.temp = as.formula(paste0(LHS, "~", RHS.mod))
  
  tf <- terms.formula(formula.temp, specials = NULL)
  trmstrings <- attr(tf, "term.labels")
  data.complete = model.frame(tf, data = data)
  
  ## construct theta matrix
  time = data.complete[time.var][,1]
  if (is.null(t.min)) {t.min = min(time)}
  if (is.null(t.max)) {t.max = max(time)}
  
  Theta = t(bs(c(t.min, t.max, time), knots = quantile(time, probs = seq(0, 1, length = Kt - 2))[-c(1,Kt - 2)], 
               intercept = TRUE, degree = 3))[,-(1:2)]
  
#  formula.model = as.formula(paste0(LHS, "~", RHS))
  formula.model = as.formula(paste0(LHS, "~ 0+", RHS))
  tf <- terms.formula(formula.model, specials = NULL)
  trmstrings <- attr(tf, "term.labels")
#  p = length(trmstrings) + 1
  p = length(trmstrings)
  mf_fixed = data.model = model.frame(tf, data = data)
  mt_fixed = attr(mf_fixed, "terms")
  
  data.model[id.var] = data.complete[id.var]
  data.model[time.var] = data.complete[time.var]
  
  ## normalize variables
  if (!standardized) {
    
    cat("Standardizing Variables \n")
    data.model = standardize_variables(data.original = data.model, time.var = time.var, trmstrings = trmstrings, LHS = LHS)
    cat("\n")
    
  }
  
  data.model[id.var] = data.complete[id.var]
  data.model[time.var] = data.complete[time.var]
  
  Y = data.model[LHS][,1]
  subj.id = data.model[id.var][,1]
  subjs = unique(subj.id)
  I = length(subjs)
  J = dim(data.model)[1]
  
  ## construct Xstar
  cat("Constructing Xstar; doing data organization \n")
  
  Xstar = tXsXs = vector("list", length = I)
  XStXS = matrix(0, Kt*p, Kt*p)
  XStY = matrix(0, Kt*p, 1)
  sumtXsXs = matrix(0, Kt*p, Kt*p)
  one.p = matrix(1, p, 1)
  one.Kt = matrix(1, Kt, 1)
  for (i in 1:I) {
    index = which(subj.id == subjs[i])
    
    X.cur = t(model.matrix(mf_fixed, data = slice(data.model, index)))
    Theta.cur = Theta[,index]
    Y.cur = Y[index]
    XS.cur = Xstar[[i]] = t(kronecker(X.cur, one.Kt) * kronecker(one.p, Theta.cur))
    XStXS = XStXS + crossprod(XS.cur)
    XStY = XStY + t(XS.cur) %*% as.matrix(Y.cur)
    sumtXsXs = sumtXsXs + t(XS.cur) %*% XS.cur
  }
  
  ## hyper parameters for inverse gaussians, bernoulli
  A = .5
  B = .5
  Atheta = .5
  Btheta = .5

  ## matrices to to approximate paramater values
  mu.q.gamma = rep(1, p)  
  mu.q.dinv = kronecker(diag((1 - mu.q.gamma)/v0 + mu.q.gamma/v1, p, p), diag(1, Kt, Kt))
  mu.q.Bpsi = matrix(0, nrow = Kt, ncol = Kp)  
  
  sigma.q.C = vector("list", I)
  for (k in 1:I) {
    sigma.q.C[[k]] = diag(1, Kp)
  }
  set.seed(1)
  mu.q.C = matrix(rnorm(I*Kp, 0, 1), I, Kp)
  
  mu.q.ltheta = digamma(Atheta + sum(mu.q.gamma)) - digamma(Atheta + Btheta + p)
  mu.q.l1.theta = digamma(Btheta + p - sum(mu.q.gamma)) - digamma(Atheta + Btheta + p)
  
  b.q.lambda.Bpsi = rep(1, Kp)
  b.q.sigma.me = 1
  
  ## initialize estimates of fixed and pca effects
  fixef.est = rep(0, J)
  pcaef.est = rep(0, J)
  
  cat("Beginning Algorithm \n")

  for (iter in 1:10) {
    
    ###############################################################
    ## update regression coefficients
    ###############################################################
    
    XStY = matrix(0, Kt*p, 1)
    for (i in 1:I) {
      index = which(subj.id == subjs[i])
      Y.cur = Y[index]
      pcaef.cur = pcaef.est[index]
      XS.cur = Xstar[[i]]
      XStY = XStY + t(XS.cur) %*% as.matrix(Y.cur - pcaef.cur)
    }
    
    sigma.q.beta = solve(as.numeric((A + J/2)/(b.q.sigma.me)) * XStXS + mu.q.dinv)
    mu.q.beta = matrix(sigma.q.beta %*% (as.numeric((A + J/2)/(b.q.sigma.me)) * XStY), nrow = Kt, ncol = p)
    
    beta.cur = t(mu.q.beta) %*% (Theta)
    
    for (i in 1:I) {
      index = which(subj.id == subjs[i])
      fixef.est[index] = Xstar[[i]] %*% as.vector(mu.q.beta)
      # equivalent formulation: apply(X[[subj]] * beta.cur, 2, sum)
    }
    
    ###############################################################
    ## update gammas
    ###############################################################
    
    for (p.cur in 1:p) {
      
      p1.div.p0 = (v1/v0) ^ (-Kt/2) * exp(-(1/2) * (1/v1 - 1/v0) * as.numeric((A + J/2)/(b.q.sigma.me)) * crossprod(mu.q.beta[,p.cur]) + (mu.q.ltheta - mu.q.l1.theta))
      
      if (p1.div.p0 > 1e50) { 
        mu.q.gamma[p.cur] = 1
      } else {
        mu.q.gamma[p.cur] = p1.div.p0/(1 + p1.div.p0)
      }
    }

    mu.q.dinv = kronecker(diag((1 - mu.q.gamma)/v0 + mu.q.gamma/v1, p, p), diag(1, Kt, Kt))
    
    mu.q.ltheta = digamma(Atheta + sum(mu.q.gamma)) - digamma(Atheta + Btheta + p)
    mu.q.l1.theta = digamma(Btheta + p - sum(mu.q.gamma)) - digamma(Atheta + Btheta + p)
    
    
    if (Kp > 0) {
      
      ###############################################################
      ## update b-spline parameters for PC basis functions
      ###############################################################
    
      sigma.q.Bpsi = solve( 
        kronecker(diag(1, Kt, Kt), diag((A + Kt/2)/b.q.lambda.Bpsi, Kp, Kp)) + 
          as.numeric((A + J/2)/(b.q.sigma.me)) * f_sum(mu.q.c = mu.q.C, sig.q.c = sigma.q.C, theta = Theta, subj.id = subj.id)
      )
      mu.q.Bpsi = matrix(((A + J/2)/(b.q.sigma.me)) * sigma.q.Bpsi %*% f_sum2(y = Y, fixef = fixef.est, subj.id = subj.id, mu.q.c = mu.q.C, kt = Kt, theta = Theta), nrow = Kt, ncol = Kp)
    
      psi.cur = t(mu.q.Bpsi) %*% (Theta)

      ###############################################################
      ## scores for each individual
      ###############################################################
    
      for (i in 1:I) {
        index = which(subj.id == subjs[i])
      
        Theta_i = Theta[,index]
        sigma.q.C[[i]] = solve( 
          diag(1, Kp, Kp ) +
            ((A + J/2)/(b.q.sigma.me)) * (f_trace(Theta_i = Theta_i, Sig_q_Bpsi = sigma.q.Bpsi, Kp = Kp, Kt = Kt) + 
            t(mu.q.Bpsi) %*% Theta_i %*% t(Theta_i) %*% mu.q.Bpsi)
        )
      
        mu.q.C[i,] = ((A + J/2)/(b.q.sigma.me)) * sigma.q.C[[i]] %*% (psi.cur[,index]) %*% (Y[index] - fixef.est[index] )
      }
    
      for (i in 1:I) {
        index = which(subj.id == subjs[i])
        pcaef.est[index] = mu.q.C[i,] %*% psi.cur[,index]
      }
    }
    
    ###############################################################
    ## update variance components
    ###############################################################
    
    if (Kp > 0) {
      ## measurement error variance
      resid = Y - fixef.est - pcaef.est
      b.q.sigma.me = as.numeric(B + .5 * (crossprod(resid) + 
                                            sum(diag(sumtXsXs %*% sigma.q.beta)) + 
                                            f_sum4(mu.q.c = mu.q.C, sig.q.c = sigma.q.C, mu.q.bpsi = mu.q.Bpsi, sig.q.bpsi = sigma.q.Bpsi, theta = Theta, subj.id = subj.id)) )
      
      ## lambda for FPCA basis functions
      for (K in 1:Kp) {
        b.q.lambda.Bpsi[K] = B + .5 * (t(mu.q.Bpsi[,K]) %*% diag(1, Kt, Kt) %*% mu.q.Bpsi[,K] + 
                                       sum(diag(diag(1, Kt, Kt) %*% sigma.q.Bpsi[(Kt*(K - 1) + 1):(Kt*K),(Kt*(K - 1) + 1):(Kt*K)])))
      }
    } else if (Kp == 0) {
      ## measurement error variance
      resid = Y - fixef.est - pcaef.est
      b.q.sigma.me = as.numeric(B + .5 * (crossprod(resid) + 
                                            sum(diag(sumtXsXs %*% sigma.q.beta)) ))
    }
    
    cat(".")
  
  }
  
  cat("\n")
  
  ## get coefficient functions over a common grid
  rownames(beta.cur) = c(trmstrings)
  beta.cur = t(beta.cur) %>% as.data.frame() %>%
    mutate(t = time) %>% 
    rename_(.dots = setNames(list("t"), time.var)) %>%
    arrange_(time.var) %>% unique() %>%
    subset(select = c(time.var, trmstrings))
  
  ## export fitted values
  Yhat = fixef.est + pcaef.est 
  
  ## export variance components
  sigeps.pm = 1 / as.numeric((A + J/2)/(b.q.sigma.me))

  ## export FPCA objects
  if (Kp > 0) {
    ## do svd to get rotated fpca basis
    rownames(psi.cur) = paste0("Psi.", 1:Kp)
    psi.cur = t(psi.cur) %>% as.data.frame() %>%
      mutate(t = time) %>% 
      arrange(t) %>% unique() %>%
      subset(select = c(paste0("Psi.", 1:Kp))) %>%
      as.matrix() %>%
      t()
    argvals = sort(unique(time))
    SVD = svd(psi.cur)
    efunctions = SVD$v
    scores = mu.q.C %*% SVD$u %*% diag(SVD$d, Kp, Kp)
    evalues = diag(cov(scores))
  } else if (Kp == 0) {
    argvals = sort(unique(time))
    efunctions = NULL
    scores = NULL
    evalues = NULL
  }
  
  Yhat.fpca = data.frame(index = time, value = pcaef.est, id = subj.id)
  class(Yhat.fpca) = c("data.frame", "refund_object")

  Y.fpca = data.frame(index = time, value = Y - fixef.est, id = subj.id)
  class(Yhat.fpca) = c("data.frame", "refund_object")
  
  fpca.obj = list(Yhat = Yhat.fpca,
                  Y = Y.fpca,
                  scores = scores,
                  mu = rep(0, length(argvals)),
                  efunctions = efunctions, 
                  evalues = evalues,
                  npc = Kp,
                  argvals = argvals)
  class(fpca.obj) = "fpca"

  ## export objects  
  ret = list(beta.cur, mu.q.gamma, mu.q.beta, Yhat, sigeps.pm, fixef.est, 
             data.complete, data.model, formula, formula.model, mt_fixed, time.var, id.var, Kt, Kp, standardized,
             t.min, t.max, fpca.obj)
  names(ret) = c("beta.pm", "gamma.pm", "spline.coef.est", "Yhat", "sigeps.pm", "fixef",
                 "data", "data.model", "formula", "formula.model", "terms", "time.var", "id.var", "Kt", "Kp", "standardized",
                 "t.min", "t.max", "fpca.obj")
  
  class(ret) = "flcm"

  ret
  
}