I. Generate Data

# Small script to simulated data that models the an idealized
# genome-wide association study, in which all single nucleotide
# polymorphisms (SNPs) have alternative alleles at high frequencies in
# the population (>5%), and all SNPs contribute some (small) amount of
# variation to the quantitative trait. The SNPs are divided into two
# sets: those that contribute a very small amount of variation, and
# those that contribute a larger amount of variable (these are the
# "QTLs"---quantitative trait loci).
#
# Adjust the five parameters below (r, d, n, p, i) to control how the
# data are enerated.
#  
# Peter Carbonetto
# University of Chicago
# January 23, 2017
#
# Diploid

# The minor alleles have additive influence on a specific trait (y).
# (certered)  X: number of minor alleles


#' @function Generate.data()
#' @param r: Proportion of variance in trait explained by SNPs.
#' @param d: Proportion of additive genetic variance due to QTLs.
#' @param n: Number of samples.
#' @param p: Number of markers (SNPs).
#' @param i: Indices of the QTLs ("causal variants").
#' @return: a list of data, with element X, y, beta1, beta2

Generate.data <- function (r = 0.6, d = 0.4, n = 400, p = 1000, i = c(27, 54)){
    # GENERATE DATA SET
    # -----------------
    # Generate the minor allele frequencies. They are uniformly
    # distributed on [0.05,0.5].
    maf <- 0.05 + 0.45 * runif(p)
    
    # Simulate genotype data X from an idealized population according to the 
    # minor allele frequencies generated above.
    X <- (runif(n*p) < maf) +
         (runif(n*p) < maf)
    X <- matrix(as.double(X), n, p, byrow = TRUE)
    
    # Center the columns of X.
    rep.row <- function (x, n)
      matrix(x, n, length(x), byrow = TRUE)
    X <- X - rep.row(colMeans(X), n)
    
    # Generate (small) polygenic additive effects for the SNPs.
    u <- rnorm(p)
    
    # Generate (large) QTL effects for the SNPs.
    beta    <- rep(0, p)
    beta[i] <- rnorm(length(i))
    
    # Adjust the additive effects so that we control for the proportion of
    # additive genetic variance that is due to QTL effects (d) and the
    # total proportion of variance explained (r). That is, we adjust beta
    # and u so that
    #
    #   r = a/(a+1)
    #   d = b/a,
    #
    # where I've defined
    #
    #   a = (u + beta)'*cov(X)*(u + beta),
    #   b = beta'*cov(X)*beta.
    #
    # Note: this code only works if d or r are not exactly 0 or exactly 1.
    st   <- c(r/(1-r) * d / var(X %*% beta))
    beta <- sqrt(st) * beta
    sa   <- max(Re(polyroot(c(c(var(X %*% beta) - r/(1-r)),
                              2*sum((X %*% beta) * (X %*% u))/n,
                              c(var(X %*% u))))))^2
    u    <- sqrt(sa) * u
    
    # Generate the quantitative trait measurements.
    y <- c(X %*% (u + beta) + rnorm(n, sd = 0.1))

    data <- list()
    data$X     <- X
    data$y     <- y
    data$beta1 <- u
    data$beta2 <- beta
    data$s1    <- sa
    data$s2    <- st 
    
    return (data)
}
# Initialize random number generation so that the results are reproducible.
set.seed(1)

r <- 0.3
d <- 0.7
n <- 50
p <- 100
idx <- c(1, 2)

data <- Generate.data(r, d, n, p, idx)
X <- data$X
y <- data$y

II. Example

Due to the numerical issues in optimizing the log-likelihood, I use a small data set where \(n=100,p=200\).

# initialize susie object s
# so that we can run update_each_effect()

init <- function (n, p, L) {
    s <- list()
    s$alpha  <- matrix(0, nrow = L, ncol = p)
    s$Xr     <- rep(0, n)
    s$mu     <- matrix(0, nrow = L, ncol = p)
    s$mu2    <- matrix(0, nrow = L, ncol = p)
    s$sigma2 <- 1
    s$pi     <- rep(1 / p, p)
    s$V      <- rep(1, L)
    s$converge <- FALSE
    
    class(s) <- "susie"
    
    return (s)
}

ridge.elbo <- function(beta.hat, X, r, res, sigma2, sigma.b2) {
    n <- nrow(X)
    p <- ncol(X)
    
    llh <- loglikelihood(sigma.b2, X, r, sigma2)
    
    return (llh + n / 2 * log (2 * pi * sigma2) + sum(res ^ 2) / (2 * sigma2))
}

ridge <- function (X, y, sigma.b2, sigma2) {
    n <- nrow(X)
    p <- ncol(X)

    S.inv <- solve(crossprod(X) + (sigma2 / sigma.b2) * diag(p))
    beta.hat <- S.inv %*% t(X) %*% y
    beta.var <- sigma2 * crossprod(X %*% S.inv)

    result <- list()
    result$beta.hat <- beta.hat
    result$var      <- beta.var

    return (result)
}

loglikelihood <- function (sigma.b2, X, y, sigma2) {
    XtX <- crossprod(X)
    Xty <- crossprod(X, y)

    p <- nrow(XtX)
    n <- nrow(X)

    S <- XtX + (sigma2 / sigma.b2) * diag(p)

    term1 <- t(Xty) %*% solve(S) %*% Xty / sigma2 - n * log(2 * pi)
    
    # sum(log(eigen(S)$values)): log(det(S))
    term2 <- - sum(log(eigen(S)$values))
    
    term3 <- - n * log (sigma2) - sum(y ^ 2) / sigma2 - p * log(sigma.b2)
    
    return ((term1 + term2 + term3) / 2)
}

ELBO.equation <- function (log.sigma.b, XtX, Xty, sigma2) {
    sigma.b2 <- exp(2 * log.sigma.b)

    p <- ncol(XtX)
    S.inv <- solve(sigma.b2 * XtX + sigma2 * diag(p))

    u <- S.inv %*% Xty

    return (sum(u ^ 2) - p)
}

MLE <- function (X, y, sigma2) {
    XtX <- crossprod(X)
    Xty <- crossprod(X, y)

    log.sigma.b <- uniroot(ELBO.equation, interval = c(-6, 6),
                           XtX = XtX, Xty = Xty, sigma2 = sigma2)$root

    return (exp(log.sigma.b * 2))
}

#'Susie mixture function
#' @param X: An n by p matrix covariates
#' @param Y: A vectore of length n response
#' @param L: Number of components
#' @param verbose: Keep track of the ELBO and print it to the console

susie.mixture <- function(X, Y, L, tol = 0.001, max_iter = 100, intercept = TRUE,
                          coverage = 0.95, min_abs_corr = 0.5, verbose = F) {
    p = ncol(X)
    n = nrow(X)
  
    s <- init(n, p, L)
    elbo <- rep(0, max_iter)
    
    ridge.bar <- rep(0, p)
      
    if (intercept) {
        mean_y <- mean(Y)
        Y <- Y - mean_y
    }  
    
    for(i in 1: max_iter){
        # susie update
        s <- update_each_effect(X, Y, s)
        
        if(verbose){
            obj <- get_objective(X, Y, s)
            if (i > 1) {
                obj <- obj + ridge.elbo(ridge.bar, X, r, ridge.res, s$sigma2, sigma.b2)
                print(paste0("objective (after susie update):", obj))
            }
        }
        
        # ridge regression update
        Y <- Y + X %*% ridge.bar
        r <- Y - s$Xr
        
        # find the optimal prior variance by maximizing the ELBO
        sigma.b2  <- MLE(X, r, s$sigma2)
        ridge.fit <- ridge(X, r, sigma.b2, s$sigma2)
        ridge.bar <- ridge.fit$beta.hat
        ridge.var <- diag(X %*% ridge.fit$var %*% t(X))
    
        Y <- Y - X %*% ridge.bar
        ridge.res <- r - X %*% ridge.bar
        
    
        if(verbose){
          
            print (mean(r^2))
            print(paste0("sigma.b^2:", sigma.b2))
            obj <- get_objective(X, Y, s) + 
                   ridge.elbo(ridge.bar, X, r, ridge.res, s$sigma2, sigma.b2)
            print(paste0("objective (after ridge update):", obj))
        }
    

        s$sigma2 = estimate_residual_variance(X, Y, s) + mean(ridge.var)
        obj <- get_objective(X, Y, s) + 
               ridge.elbo(ridge.bar, X, r, ridge.res, s$sigma2, sigma.b2)
        
        if(verbose)
            print(paste0("objective (after update sigma2):", obj))

        elbo[i + 1] <- obj
        if((elbo[i + 1] - elbo[i]) < tol) {
            s$converged <- TRUE
            break;
        }
    }
    
    elbo    <- elbo[2: (i + 1)] # Remove first (infinite) entry, and trailing NAs.
    s$elbo  <- elbo
    s$niter <- i
    
    if (is.null(s$converged)) {
        warning(paste("IBSS algorithm did not converge in", max_iter, "iterations!"))
        s$converged = FALSE
    }
    
    if(intercept){
        s$intercept = mean_y - sum(attr(X,"scaled:center") *  (colSums(s$alpha*s$mu)/attr(X,"scaled:scale")))# estimate intercept (unshrunk)
        s$fitted = s$Xr + mean_y
    } else {
        s$intercept <- 0
        s$fitted <- s$Xr
    }
    s$fitted <- drop(s$fitted)

    
    ## SuSiE PIP
    s$null_index = 0
    s$pip = susie_get_pip(s)

    ## mixture
    s$sigma.b2   <- sigma.b2
    s$ridge.coef <- ridge.fit$beta.hat

    return(s)
}


library(R.utils)

sourceDirectory("R_function")

s  <- susie(X, y, L = 2, verbose = T, scaled_prior_variance = 0.5,
            estimate_prior_variance = FALSE, standardize = FALSE) 

[1] "objective:-44.5319038772319"
[1] "objective:-38.7993191041671"
[1] "objective:-36.2925884480723"
[1] "objective:-34.9486493391848"
[1] "objective:-34.647911036539"
[1] "objective:-34.6204326163046"
[1] "objective:-34.6177278158765"
[1] "objective:-34.617616355774"
[1] "objective:-34.6176066780397"
[1] "objective:-34.6176062968079"

sm <- susie.mixture(X, y, L = 2, verbose = T)

[1] 0.3978632
[1] "sigma.b^2:0.0438337538720929"
[1] "objective (after ridge update):-76.5966437053549"
[1] "objective (after update sigma2):-32.784976098098"

par(mfrow = c(2, 2))

plot(susie_get_pip(sm), abs(data$beta1 + data$beta2), 
     main = "SuSiE mixture", xlab = "PIP", ylab = "Random effect + Fixed effect")
plot(susie_get_pip(sm), abs(data$beta2), 
     main = "SuSiE mixture", xlab = "PIP", ylab = "Only fixed effect")
plot(susie_get_pip(s), abs(data$beta1 + data$beta2), 
     main = "SuSiE", xlab = "PIP", ylab = "Random effect + Fixed effect")
plot(susie_get_pip(s), abs(data$beta2), 
     main = "SuSiE", xlab = "PIP", ylab = "Only fixed effect")

Past versions of unnamed-chunk-3-1.png

As expected, SuSiE-mixture model shrinks the PIP, and makes less discovery.

sessionInfo()

R version 3.6.0 (2019-04-26)
Platform: x86_64-w64-mingw32/x64 (64-bit)
Running under: Windows 10 x64 (build 17134)

Matrix products: default

locale:
[1] LC_COLLATE=English_United States.1252 
[2] LC_CTYPE=English_United States.1252   
[3] LC_MONETARY=English_United States.1252
[4] LC_NUMERIC=C                          
[5] LC_TIME=English_United States.1252    

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

other attached packages:
[1] R.utils_2.9.0     R.oo_1.22.0       R.methodsS3_1.7.1

loaded via a namespace (and not attached):
 [1] workflowr_1.4.0 Rcpp_1.0.1      lattice_0.20-38 digest_0.6.19  
 [5] rprojroot_1.3-2 grid_3.6.0      backports_1.1.4 git2r_0.25.2   
 [9] magrittr_1.5    evaluate_0.14   stringi_1.4.3   fs_1.3.1       
[13] whisker_0.3-2   Matrix_1.2-17   rmarkdown_1.13  tools_3.6.0    
[17] stringr_1.4.0   glue_1.3.1      xfun_0.7        yaml_2.2.0     
[21] compiler_3.6.0  htmltools_0.3.6 knitr_1.23