R/IBSS.R

In statwangz/SuSiE: A simple new approach to variable selection in regression, with application to genetic fine-mapping

#' Iterative Bayesian stepwise selection
#'
#' @param X SNPs, a n*p matrix.
#' @param y trait, a n*1 vector.
#' @param L number of the compoments of the effect variables.
#' @param pi0 hyperparameter of the prior distribution of gamma.
#' @param sig2 hyperparameter, the variance of the trait y.
#' @param sig02 hyperparameter, the variance of the coefficient.
#' @param tol convergence criterion。
#' @param iter_max the max number of the iteration.
#'
#' @return the SER result of each compoment and the PIP in each iteration.
#' @export
IBSS <- function(X, y, L = 3, pi0 = NULL, sig2 = 1, sig02 = 0.1, tol = 1e-5, iter_max = 1000) {
  n <- length(y)
  p <- ncol(X)

  if (is.null(pi0)) {
    pi0 <- rep(1 / p, p)
  }

  iter <- 0 # Count
  ser <- list() # Store the result of the function SER()

  # Store the marginal posterior inclusion probability (PIP) of p variables
  pip <- tibble(idx = paste0("X", 1:p))
  pip_names <- paste0("alpha", 1:L)
  pip_set <- matrix(0, nrow = p, ncol = L) %>%
    as_tibble(.name_repair = "minimal") %>%
    set_names(pip_names) %>%
    bind_cols(tibble(idx = paste0("X", 1:p)), .)

  b_new <- matrix(0, nrow = p, ncol = L) # Initialize b

  repeat{
    b_old <- b_new
    iter <- iter + 1

    for (l in 1:L) {
      r <- y - rowSums(X %*% b_new[, -l])
      ser[[l]] <- SER(X, r, pi0, sig2, sig02)
      b_new[, l] <- ser[[l]]$alpha * ser[[l]]$mu1
      pip_set[, l + 1] <- ser[[l]]$alpha
    }

    # Compute the PIP: PIP_j = 1 - prod(1 - alpha_lj, l from 1 to L)
    pip_set %>%
      pivot_longer(all_of(pip_names),
        names_to = "set/l",
        values_to = "alpha"
      ) %>%
      group_by(idx) %>%
      mutate(PIP = 1 - prod(1 - alpha)) %>%
      ungroup() %>%
      pivot_wider(
        names_from = "set/l",
        values_from = "alpha"
      ) %>%
      select(PIP) %>%
      set_names(paste0("iter", iter)) %>%
      bind_cols(pip, .) -> pip

    if (norm(b_new - b_old, type = "F") < tol || iter == iter_max) {
      break
    }
  }

  return(ibss = list(ser = ser, PIP = pip))
}