R/StepTwo.R

In SMAHP: Survival Mediation Analysis of High-Dimensional Proteogenomic Data

#' Select M-X pairs via SIS
#'
#' An auxiliary function which selects M-X paris using SIS.
#' @param X An n by p matrix of exposures.
#' @param M An n by p matrix of mediators.
#' @param time A vector of survival time of samples.
#' @param status A vector of status indicator: 0=alive, 1=dead.
#' @param X_sel_Y_s1 Outputs from StepOne: A vector of selected X for Y.
#' @param M_X_s1 Outputs from StepOne: A list of selected X for M.
#' @param M_sel_Y_s1 Outputs from StepOne: A vector of selected M for Y.
#' @param SIS_thres SIS thresholds. Default is "n/log(n)". Other options include "n-1", "n/2log(n)", "2n/log(n)", "3n/log(n)".
#' @return A data table with selected M, X pairs and related effect size.
#' @import data.table
#' @import survival
#' @export
#' @examples
#' \donttest{
#' data(example_dat)
#' surv_dat <- example_dat$surv_dat
#' res_step1 <- StepOne(X = example_dat$X, M = example_dat$M, time = surv_dat$time,
#' status = surv_dat$status, model_option = "MCP")
#' M_X_sel_s2 <- StepTwo(X = example_dat$X, M = example_dat$M, time = surv_dat$time,
#' status = surv_dat$status, X_sel_Y_s1 = res_step1$X_sel_Y_s1, M_X_s1 = res_step1$M_X_s1,
#' M_sel_Y_s1 = res_step1$M_sel_Y_s1)
#' }

StepTwo <- function(X, M, time, status, X_sel_Y_s1, M_X_s1, M_sel_Y_s1, SIS_thres = "n/log(n)") {
  n <- nrow(X)
  beta_m_s2 <- c()
  surv_data <- data.frame(surv_time = time, status = status)
  for(m in M_sel_Y_s1){
    surv_model_data_s2 <- cbind(M[ ,m], X[ ,X_sel_Y_s1], surv_data)
    surv_model_s2 <- survreg(Surv(surv_time, status) ~ ., data = surv_model_data_s2, dist = "lognormal")
    beta_m_s2 <- c(beta_m_s2, as.numeric(surv_model_s2$coefficients[2]))
  }
  names(beta_m_s2) <- M_sel_Y_s1

  M_X_s2 <- list()
  for(m in M_sel_Y_s1){
    X_sel <- M_X_s1[[m]]
    lm_model_data_s2 <- cbind(M = M[ ,m], X = X[ ,X_sel])
    lm_model_data_s2 <- as.data.table(lm_model_data_s2)
    lm_model_s2 <- lm(M ~ ., data = lm_model_data_s2)
    lm_model_s2_results <- summary(lm_model_s2)$coefficients[ ,1]
    if(length(lm_model_s2_results) == 1){
      M_X_s2[[m]] <- NULL
      next
    }
    alpha <- lm_model_s2$coefficients[-1]
    if(length(X_sel) == 1){
      names(alpha) <- X_sel
    }
    beta <- beta_m_s2[m]
    lm_model_s2_results <- data.table(M = m, X = names(alpha), alpha = as.numeric(alpha), beta = beta)
    #lm_model_s2_results[ , effect := abs(beta*alpha)]
    lm_model_s2_results$effect <- abs(lm_model_s2_results$beta * lm_model_s2_results$alpha)
    M_X_s2[[m]] <- lm_model_s2_results
  }
  M_X_s2 <- rbindlist(M_X_s2)
  #M_X_s2_sorted <- M_X_s2[order(effect, decreasing = TRUE)]
  M_X_s2_sorted <- M_X_s2[order(M_X_s2$effect, decreasing = TRUE), ]

  if(SIS_thres == "n/log(n)"){
    d_s2 <- n/log(n)
  }
  if(SIS_thres == "2n/log(n)"){
    d_s2 <- 2*n/log(n)
  }
  if(SIS_thres == "3n/log(n)"){
    d_s2 <- 3*n/log(n)
  }
  if(SIS_thres == "n-1"){
    d_s2 <- n - 1
  }
  if(SIS_thres == "n/2log(n)"){
    d_s2 <- n/(2*log(n))
  }

  if(d_s2 <= nrow(M_X_s2_sorted)){
    M_X_sel_s2 <- M_X_s2_sorted[1:d_s2, ]
  } else {
    M_X_sel_s2 <- M_X_s2_sorted
  }

  return(M_X_sel_s2)

}