Revisiting prior families for trees: Part IV

workflowr

Ignored files:
    Ignored:    .DS_Store
    Ignored:    .Rhistory
    Ignored:    .Rproj.user/
    Ignored:    docs/.DS_Store
    Ignored:    docs/assets/.DS_Store
    Ignored:    output/

Untracked files:
    Untracked:  analysis/extrapolate3.Rmd
    Untracked:  analysis/extrapolate4.Rmd

Unstaged changes:
    Modified:   drift-workflow.Rproj

suppressMessages({
  library(flashier)
  library(drift.alpha)
  library(tidyverse)
})

I use the more challenging scenario from the previous analysis, but I’d like to see whether I can simplify the fitting process. Namely, I’ll use a mixture of point masses at -1, 0, and +1 without fixing any loadings at zero. Hopefully, I’ll be able to get a tree-like structure without imposing it via external constraints. I’ll also attempt to do the relaxation without fixed loadings. To do so, I’ll need to keep all three point masses as well as adding a uniform component on \([-0.9, 0.9]\) (in the previous analysis, I removed the point mass at zero).

set.seed(666)

nA <- 100
nB <- 50
nC <- 40
nD <- 80
nE <- 50
n <- nA + nB + nC + nD + nE

p <- 10000

a <- rnorm(p, sd = 2)
b <- rnorm(p, sd = 1.5)
c <- rnorm(p)
d <- rnorm(p, sd = 0.5)
e <- rnorm(p, sd = 0.25)
f <- rnorm(p, sd = 0.5)
g <- rnorm(p, sd = 0.75)

popA <- c(rep(1, nA), rep(0, nB + nC + nD + nE))
popB <- c(rep(0, nA), rep(1, nB), rep(0, nC + nD + nE))
popC <- c(rep(0, nA + nB), rep(1, nC), rep(0, nD + nE))
popD <- c(rep(0, nA + nB + nC), rep(1, nD), rep(0, nE))
popE <- c(rep(0, nA + nB + nC + nD), rep(1, nE))

E.factor <- 0.10 * (a + b + e) + 0.35 * (a + c + f) + 0.55 * (a + c + g)

Y <- cbind(popA, popB, popC, popD, popE) %*% 
  rbind(a + b + d, a + b + e, a + c + f, a + c + g, E.factor)
Y <- Y + rnorm(n * p, sd = 0.1)

plot_dr <- function(dr) {
  sd <- sqrt(dr$prior_s2)
  L <- dr$EL
  LDsqrt <- L %*% diag(sd)
  K <- ncol(LDsqrt)
  plot_loadings(LDsqrt[,1:K], rep(c("A", "B", "C", "D", "E"), c(nA, nB, nC, nD, nE))) +
    scale_color_brewer(palette="Set2")
}

tree.fn = function(x, s, g_init, fix_g, output, ...) {
  if (is.null(g_init)) {
    g_init <- ashr::unimix(rep(1/3, 3), c(-1, 0, 1), c(-1, 0, 1))
  }

  return(flashier:::ebnm.nowarn(x = x,
                                s = s,
                                g_init = g_init,
                                fix_g = fix_g,
                                output = output,
                                prior_family = "ash",
                                prior = c(10, 1, 10),
                                ...))
}

prior.tree = function(...) {
  return(as.prior(sign = 0, ebnm.fn = function(x, s, g_init, fix_g, output) {
    tree.fn(x, s, g_init, fix_g, output, ...)
  }))
}

flextree.fn = function(x, s, g_init, fix_g, output, ...) {
  if (is.null(g_init)) {
    g_init <- ashr::unimix(rep(1/4, 4), c(-1, 0, 1, -0.9), c(-1, 0, 1, 0.9))
  }

  return(flashier:::ebnm.nowarn(x = x,
                                s = s,
                                g_init = g_init,
                                fix_g = fix_g,
                                output = output,
                                prior_family = "ash",
                                prior = c(10, 10, 10, 1),
                                ...))
}

prior.flextree = function(...) {
  return(as.prior(sign = 0, ebnm.fn = function(x, s, g_init, fix_g, output) {
    flextree.fn(x, s, g_init, fix_g, output, ...)
  }))
}

init.mean.factor <- function(resids, zero.idx) {
  u <- matrix(1, nrow = nrow(resids), ncol = 1)
  u[zero.idx, 1] <- 0
  v <- t(solve(crossprod(u), crossprod(u, resids)))
  return(list(u, v))
}

init.split.factor <- function(resids, zero.idx) {
  svd.res <- svd(resids, nu = 1, nv = 1)
  u <- svd.res$u
  u[zero.idx] <- 0
  u <- matrix(sign(u), ncol = 1)
  v <- t(solve(crossprod(u), crossprod(u, resids)))
  return(list(u, v))
}

If I use the correct value of K, the initial tree appears as follows:

fl <- flash.init(Y) %>%
  flash.set.verbose(0) %>%
  flash.init.factors(EF = init.mean.factor(Y, NULL), 
                     prior.family = c(prior.tree(), prior.normal())) %>%
  flash.fix.loadings(kset = 1, mode = 1L) %>%
  flash.backfit() %>%
  flash.add.greedy(Kmax = 3, prior.family = c(prior.tree(), prior.normal())) %>%
  flash.backfit()

plot_dr(init_from_flash(fl))

Not bad! I’ll leave the question of misspecified K aside for now (and indeed, it might not be a big issue in practice) and see whether the relaxation works without fixing loadings.

fl2 <- fl

# Relax the priors.
for (k in 1:4) {
  fl2$flash.fit$ebnm.fn[[k]][[1]] <- flextree.fn
}
fl2 <- fl2 %>% flash.backfit(warmstart = FALSE) %>% flash.backfit()

plot_dr(init_from_flash(fl2))

It does! As in the previous analysis, the loadings for Population \(E\) are correct:

LL <- colMeans(fl2$flash.fit$EF[[1]][271:320, ])
names(LL) <- paste("Factor", 1:4)
round(LL, 2)

#> Factor 1 Factor 2 Factor 3 Factor 4 
#>      1.0     -0.8     -0.2      0.1

Session information

sessionInfo()

#> R version 3.5.3 (2019-03-11)
#> Platform: x86_64-apple-darwin15.6.0 (64-bit)
#> Running under: macOS Mojave 10.14.6
#> 
#> Matrix products: default
#> BLAS: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRblas.0.dylib
#> LAPACK: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRlapack.dylib
#> 
#> locale:
#> [1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8
#> 
#> attached base packages:
#> [1] stats     graphics  grDevices utils     datasets  methods   base     
#> 
#> other attached packages:
#>  [1] forcats_0.4.0     stringr_1.4.0     dplyr_0.8.0.1    
#>  [4] purrr_0.3.2       readr_1.3.1       tidyr_0.8.3      
#>  [7] tibble_2.1.1      ggplot2_3.2.0     tidyverse_1.2.1  
#> [10] drift.alpha_0.0.9 flashier_0.2.5   
#> 
#> loaded via a namespace (and not attached):
#>  [1] Rcpp_1.0.4.6       lubridate_1.7.4    invgamma_1.1      
#>  [4] lattice_0.20-38    assertthat_0.2.1   rprojroot_1.3-2   
#>  [7] digest_0.6.18      truncnorm_1.0-8    R6_2.4.0          
#> [10] cellranger_1.1.0   plyr_1.8.4         backports_1.1.3   
#> [13] evaluate_0.13      httr_1.4.0         pillar_1.3.1      
#> [16] rlang_0.4.2        lazyeval_0.2.2     readxl_1.3.1      
#> [19] rstudioapi_0.10    ebnm_0.1-21        irlba_2.3.3       
#> [22] whisker_0.3-2      Matrix_1.2-15      rmarkdown_1.12    
#> [25] labeling_0.3       munsell_0.5.0      mixsqp_0.3-40     
#> [28] broom_0.5.1        compiler_3.5.3     modelr_0.1.5      
#> [31] xfun_0.6           pkgconfig_2.0.2    SQUAREM_2017.10-1 
#> [34] htmltools_0.3.6    tidyselect_0.2.5   workflowr_1.2.0   
#> [37] withr_2.1.2        crayon_1.3.4       grid_3.5.3        
#> [40] nlme_3.1-137       jsonlite_1.6       gtable_0.3.0      
#> [43] git2r_0.25.2       magrittr_1.5       scales_1.0.0      
#> [46] cli_1.1.0          stringi_1.4.3      reshape2_1.4.3    
#> [49] fs_1.2.7           xml2_1.2.0         generics_0.0.2    
#> [52] RColorBrewer_1.1-2 tools_3.5.3        glue_1.3.1        
#> [55] hms_0.4.2          parallel_3.5.3     yaml_2.2.0        
#> [58] colorspace_1.4-1   ashr_2.2-51        rvest_0.3.4       
#> [61] knitr_1.22         haven_2.1.1