Revisiting prior families for trees: Part IX

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
    Untracked:  analysis/pm1_priors10.Rmd

Unstaged changes:
    Modified:   drift-workflow.Rproj

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

I return to the simplest scenario (a balanced tree with four populations of equal sizes) to see whether we can get the tree using a less contrived family of priors than the three-pointmass family.

sim_tree <- function(n_range, 
                     p = 10000, 
                     branch_means, 
                     branch_sds, 
                     resid_sd = 0.1,
                     admix_pops = NULL, 
                     outgroup = FALSE, 
                     seed = 666) {
  set.seed(666)
  
  depth <- length(branch_means)
  npop_pure <- 2^(depth - 1)

  if (is.null(admix_pops)) {
    admix_pops <- matrix(nrow = 0, ncol = 0)
  }
  npop_admix <- ncol(admix_pops)
  
  npop <- npop_pure + npop_admix + outgroup
  
  if (length(n_range) == 1) {
    n <- rep(n_range, npop)
  } else {
    n <- sample(30:100, npop, replace = TRUE)
  }
  K <- 2^depth - 1
  
  FF <- matrix(nrow = p, ncol = K)
  k <- 1
  for (d in 1:depth) {
    for (i in 1:(2^(d - 1))) {
      FF[, k] <- rnorm(p, sd = branch_means[d] + rnorm(1, sd = branch_sds[d]))
      k <- k + 1
    }
  }
  
  tree_mat <- matrix(0, nrow = npop_pure, ncol = K)
  k <- 1
  for (d in 1:depth) {
    size <- 2^(depth - d)
    for (i in 1:(2^(d - 1))) {
      tree_mat[((i - 1) * size + 1):(i * size), k] <- 1
      k <- k + 1
    }
  }
  
  pop_means <- FF %*% t(tree_mat)
  if (npop_admix > 0) {
    pop_means <- cbind(pop_means, pop_means %*% admix_pops)
  }
  if (outgroup) {
    pop_means <- cbind(pop_means, rnorm(p, mean = 0, sd = sqrt(sum(branch_sds^2))))
  }
  
  Y <- NULL
  for (i in 1:npop) {
    Y <- rbind(Y, matrix(pop_means[, i], nrow = n[i], ncol = p, byrow = TRUE))
  }
  Y <- Y + rnorm(sum(n) * p, sd = resid_sd)
  
  plot_fl <- function(fl) {
    dr <- init_from_flash(fl)
    sd <- sqrt(dr$prior_s2)
    L <- dr$EL
    LDsqrt <- L %*% diag(sd)
    K <- ncol(LDsqrt)
    plot_loadings(LDsqrt[,1:K], rep(letters[1:npop], n)) +
      scale_color_brewer(palette="Set3")
  }
  
  return(list(Y = Y, plot_fn = plot_fl))
}
  
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))
}

Normal mixture priors fail to find the correct rotation:

balanced_4pop <- sim_tree(n_range = 50,
                          p = 10000,
                          branch_means = rep(1, 3),
                          branch_sds = rep(0, 3),
                          resid_sd = 0.1)

fl_nsm <- flash.init(balanced_4pop$Y) %>%
  flash.set.verbose(0) %>%
  flash.add.greedy(Kmax = 4, 
                   prior.family = c(prior.normal.scale.mix(), prior.normal())) %>%
  flash.backfit(tol = 1e-4)

balanced_4pop$plot_fn(fl_nsm)

Past versions of normal_mix-1.png

Unimodal priors can’t find it either:

fl_uni <- flash.init(balanced_4pop$Y) %>%
  flash.set.verbose(0) %>%
  flash.add.greedy(Kmax = 4, 
                   prior.family = c(prior.unimodal(), prior.normal())) %>%
  flash.backfit(tol = 1e-4, maxiter = 1000, verbose.lvl = 0)

#> Warning in report.maxiter.reached(verbose.lvl): Maximum number of
#> iterations reached.

balanced_4pop$plot_fn(fl_uni)

Past versions of unimix-1.png

Even varimax rotations won’t find it:

add_split <- function(fl, Y, eigen.thresh = 0.5) {
  svd.res <- svd(Y - fitted(fl), nu = 10, nv = 10)
  n.eigen <- sum(svd.res$d > eigen.thresh * svd.res$d[1])
  if (n.eigen > 1) {
    varimax.res <- varimax(svd.res$u[, 1:n.eigen], normalize = FALSE)
    EL <- varimax.res$loadings[, 1, drop = FALSE]
    EF <- (svd.res$v[, 1:n.eigen] %*% t(solve(varimax.res$rotmat)))[, 1, drop = FALSE]
  } else {
    EL <- svd.res$u[, 1, drop = FALSE]
    EF <- svd.res$v[, 1, drop = FALSE]
  }
  fl.new <- fl %>%
    flash.init.factors(EF = list(EL, EF),
                       prior.family = c(prior.normal.scale.mix(), prior.normal())) %>%
    flash.backfit(kset = fl$n.factors + 1)
  return(fl.new)
}

fl_varimax <- flash.init(balanced_4pop$Y) %>%
  flash.set.verbose(0) %>%
  flash.add.greedy(Kmax = 2, 
                   prior.family = c(prior.normal.scale.mix(), prior.normal()))
for (i in 1:2) {
  fl_varimax <- add_split(fl_varimax, balanced_4pop$Y)
}
fl_varimax <- flash.backfit(fl_varimax, tol = 1e-4)

balanced_4pop$plot_fn(fl_varimax)

Point-Laplace priors, however, do find the tree!

fl_pl <- flash.init(balanced_4pop$Y) %>%
  flash.set.verbose(0) %>%
  flash.add.greedy(Kmax = 4, 
                   prior.family = c(prior.point.laplace(), prior.normal())) %>%
  flash.backfit(tol = 1e-4, verbose.lvl = 0)

balanced_4pop$plot_fn(fl_pl)

Past versions of pl-1.png

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.10 flashier_0.2.7    
#> 
#> 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