SBC_DESeq2.Rmd

---
title: "SBC & DESeq2"
date: 2021-09-23
output: html_document
---
  
Let's setup the environment.

```{r setup, message=FALSE,warning=FALSE, results="hide"}
library(SBC)
library(ggplot2)

# Bioconductor packages
library(DESeq2) 
library(apeglm)

# Setup multiprocessing
library(future)
plan(tweak(multisession, workers = 16))
#plan(multisession)

theme_set(cowplot::theme_cowplot())

# Setup caching of results
cache_dir <- "./_sbc_deseq2_cache/"
if(!dir.exists(cache_dir)) {
  dir.create(cache_dir)
}

image_path <- "."
image_width <- 5
image_height <- 3

```

Define an SBC backend using DESeq2 to sample the regression coefficients.

```{r}
SBC_backend_DESeq2 <- function(shrink_type = "normal", ..., useT = FALSE) {
    shrink_args = list(...)
    if(any(names(shrink_args) %in% c("coef", "contrast", "type"))) {
      #TODO better report error
      stop(paste0("Parameter 'TODO' cannot be provided when defining DESeq2 backend, must be passed other way"))
    }
  
    structure(list(shrink_type = shrink_type, 
                   shrink_args = shrink_args,
                   useT = useT), class = "SBC_backend_DEseq2")
}

SBC_fit.SBC_backend_DEseq2 <- function(backend, generated, cores) {
  colData <- data.frame(group = factor(rep(c("A", "B"), each = generated$N_per_group)))
  dds <- DESeq2::DESeqDataSetFromMatrix(countData = generated$countData,
                              colData = colData,
                              design = ~ group)
  
  DESeq2::sizeFactors(dds) <- rep(1, generated$N_per_group * 2)
  
  dds <- DESeq2::DESeq(dds, fitType = "parametric", 
                       minReplicatesForReplace = Inf, minmu = 1e-10,
                       useT = backend$useT)
  
#  if(backend$shrink_type != "none") {
  # } else {
  #   res <- DESeq2::results()
  # }

  structure( list(
    dds = dds,
    backend = backend),
    class = "SBC_fit_DESeq2"
  )  
  
}

SBC_fit_to_draws_matrix.SBC_fit_DESeq2 <- function(fit) {
  if(fit$backend$shrink_type != "none") {
    res <- do.call(DESeq2::lfcShrink,
                 c(list(dds = fit$dds, coef = "group_B_vs_A", 
                        type = fit$backend$shrink_type),
                   fit$backend$shrink_args))
  } else {
    res <- DESeq2::results(fit$dds, contrast = c("group", "B", "A"))
  }

  n_samps <- 1000
  
  useT <- "tDegreesFreedom" %in% names(S4Vectors::mcols(fit$dds))
  if(useT) {
    df <- unique(S4Vectors::mcols(fit$dds)$tDegreesFreedom)
    if(length(df) > 1) {
      stop("Non-unique DF")
    }
    base_samps <- rt(n_samps * nrow(res), df = df)
  } else {
    base_samps <- rnorm(n_samps * nrow(res))
  }
  beta_samps <- matrix(res$log2FoldChange, nrow = n_samps, ncol = nrow(res),  byrow = TRUE) +
      matrix(base_samps, nrow = n_samps, ncol = nrow(res)) *
       matrix(res$lfcSE, nrow = n_samps, ncol = nrow(res),  byrow = TRUE)

  posterior::as_draws_matrix(posterior::draws_rvars(beta = posterior::rvar(beta_samps)))
}
```





```{r}
SBC_backend_iid_draws.SBC_backend_DEseq2 <- function(backend) {
  TRUE
}
```


Code to simulate data from the DESeq2 model

```{r}
generator_single_DEseq2 <- function(N_per_group,
                                    N_genes, 
                                    intercept_prior_mean = 4, 
                                    intercept_prior_sd = 2, 
                                    predictor_prior_sd = 1) {
  
  asymptDisp <- rgamma(1, shape = 3, rate = 6)
  extraPois <- rgamma(1, shape = 4, rate = 2.3)
  disp_sd = abs(rnorm(1,0,1))

  countData <- matrix(NA_integer_, nrow = N_genes, ncol = 2 * N_per_group)
  X <- rep(c(0, 1), each = N_per_group)
  
  
  # Do rejection sampling to avoid all zero rows
  needs_recomputing <- rep(TRUE, N_genes)

  intercept <- array(NA_real_, N_genes)
  beta <- array(NA_real_, N_genes)

  repeat {   
    N_recomputing <- sum(needs_recomputing)
    intercept[needs_recomputing] <- rnorm(N_recomputing, intercept_prior_mean, intercept_prior_sd)
    beta[needs_recomputing] <- rnorm(N_recomputing, 0, predictor_prior_sd)

    mean_expr <- 2^intercept + 0.5 * (2^beta)
    disp_mean <-  extraPois/mean_expr + asymptDisp
    dispersions <- rlnorm(N_recomputing, disp_mean, disp_sd)

    for(samp in 1:(2 * N_per_group)) {
      logmus <- intercept[needs_recomputing] + beta[needs_recomputing] * X[samp]
      mus <- 2^logmus
      countData[needs_recomputing, samp] <- rnbinom(N_recomputing, mu = mus, size = 1/dispersions)
    }
    needs_recomputing <- rowSums(countData) == 0
    if(!any(needs_recomputing)) {
      break;
    }
  }
  
  rownames(countData) <- paste0("g", 1:N_genes)
  
  list(
    variables = list(beta = beta),
    generated = list(N_per_group = N_per_group, countData = countData)
  ) 
}
```


```{r}
merge_stats <- function(res) {
  merged_stats <- dplyr::mutate(res$stats, variable = "Merged", sim_id = 1:dplyr::n())
  merged_stats <- dplyr::filter(merged_stats, !is.na(rank))
  merged_stats
}

overall_coverage <- function(res, width = 0.95) {
  merged_stats <- merge_stats(res)
  #dplyr::select(empirical_coverage(merged_stats, width), width, estimate)
  empirical_coverage(merged_stats, width)
}

plot_coverage_diff_merged <- function(res) {
  plot_coverage_diff(merge_stats(res), max_points = 101)  +
    geom_vline(color = "orangered", alpha = 0.8, xintercept = 0.95)
}
```


```{r}
set.seed(43216842)
datasets20 <- generate_datasets(SBC_generator_function(
   generator_single_DEseq2, N_per_group = 20, N_genes = 1000),
   n_sims = 1000) 
```

## 20 samples, normal

```{r}
res20 <- compute_SBC(datasets20, SBC_backend_DESeq2(),
                         keep_fits = FALSE,
                         globals = c("SBC_fit.SBC_backend_DEseq2", 
                                     "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                     "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_20"))
```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res20, variables = pars_to_show)
plot_ecdf_diff(res20, variables = pars_to_show)
plot_sim_estimated(res20, variables = pars_to_show)
plot_coverage_diff(res20, variables = pars_to_show)
```

```{r}
overall_coverage(res20)
```


## 3 samples, normal

```{r}
set.seed(26845214)
datasets3 <- generate_datasets(SBC_generator_function(
   generator_single_DEseq2, N_per_group = 3, N_genes = 1000),
   n_sims = 1000) 
```


```{r}
res3 <- compute_SBC(datasets3, SBC_backend_DESeq2(), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3"),
                        keep_fits = FALSE)
```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res3, variables = pars_to_show)
plot_ecdf_diff(res3, variables = pars_to_show)
plot_sim_estimated(res3, variables = pars_to_show)
plot_coverage_diff(res3, variables = pars_to_show)
```

```{r}
overall_coverage(res3)
```

```{r}
p <- plot_coverage_diff_merged(res3)
p
ggsave(file.path(image_path,"3_normal_diff.png"), p,
  width = image_width, height = image_height)
```


## 3 samples, normal, useT

```{r}
res3_t <- compute_SBC(datasets3, SBC_backend_DESeq2(useT = TRUE), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3_T"),
                        keep_fits = FALSE)
```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res3_t, variables = pars_to_show)
plot_ecdf_diff(res3_t, variables = pars_to_show)
plot_sim_estimated(res3_t, variables = pars_to_show)
plot_coverage_diff(res3_t, variables = pars_to_show)
```

```{r}
overall_coverage(res3_t)
```

```{r}
p <- plot_coverage_diff_merged(res3_t)
p
ggsave(file.path(image_path,"3_normal_t_diff.png"), p,
  width = image_width, height = image_height)
```


## 3 samples, none

```{r}
res3_none <- compute_SBC(datasets3, SBC_backend_DESeq2(shrink_type = "none"), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3_none"),
                        keep_fits = FALSE)

```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res3_none, variables = pars_to_show)
plot_ecdf_diff(res3_none, variables = pars_to_show)
plot_sim_estimated(res3_none, variables = pars_to_show)
plot_coverage_diff(res3_none, variables = pars_to_show)
```

```{r}
overall_coverage(res3_none)
```

```{r}
p <-  plot_coverage_diff_merged(res3_none)
p
ggsave(file.path(image_path,"3_none_diff.png"),
  p,
  width = image_width, height = image_height)
```

```{r}
empirical_coverage(dplyr::mutate(dplyr::filter(res3_none$stats, q5 > 0 | q95 < 0), variable = "Merged"), width = 0.95)
```

## 3 samples, none, use T

```{r}
res3_none_t <- compute_SBC(datasets3, SBC_backend_DESeq2(shrink_type = "none", useT = TRUE), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3_none_t"),
                        keep_fits = FALSE)
```

```{r}
overall_coverage(res3_none_t)
```


```{r}
p <- plot_coverage_diff_merged(res3_none_t)
p
ggsave(file.path(image_path,"3_none_t_diff.png"), p,
  width = image_width, height = image_height)
```

## 20 samples, apeglm

```{r}
res20_apeglm <- compute_SBC(datasets20, SBC_backend_DESeq2(shrink_type = "apeglm"), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_20_apeglm"),
                        keep_fits = FALSE)
```

```{r}
save_dir <- "_sbc_deseq2_cache/"

plot_rank_hist(res20_apeglm, variables = "beta[109]")
plot_ecdf_diff(res20_apeglm, variables = "beta[109]")
plot_coverage_diff(res20_apeglm, variables = "beta[109]")

```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res20_apeglm, variables = pars_to_show)
plot_ecdf_diff(res20_apeglm, variables = pars_to_show)
plot_sim_estimated(res20_apeglm, variables = pars_to_show, alpha = 0.1)
plot_coverage_diff(res20_apeglm, variables = pars_to_show)

```
```{r}
overall_coverage(res20_apeglm)
```

```{r}
p <-  plot_coverage_diff_merged(res20_apeglm)
p
ggsave(file.path(image_path,"20_apeglm_diff.png"),
  p,
  width = image_width, height = image_height)
```


## 20 samples, apeglm, useT

```{r}
res20_apeglm_T <- compute_SBC(datasets20, SBC_backend_DESeq2(shrink_type = "apeglm", useT = TRUE), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_20_apeglm_T"),
                        keep_fits = FALSE)
```



```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res20_apeglm_T, variables = pars_to_show)
plot_ecdf_diff(res20_apeglm_T, variables = pars_to_show)
plot_sim_estimated(res20_apeglm_T, variables = pars_to_show, alpha = 0.1)
plot_coverage_diff(res20_apeglm_T, variables = pars_to_show)

```
```{r}
overall_coverage(res20_apeglm_T, width = c(0.9,0.95,0.98))
```

## 3 samples, apeglm

```{r}
res3_apeglm <- compute_SBC(datasets3, SBC_backend_DESeq2(shrink_type = "apeglm"), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3_apeglm"),
                        keep_fits = FALSE)
```
```{r}
rank_hist_1 <- plot_rank_hist(res3_apeglm, variables = "beta[745]")
plot_ecdf_diff(res3_apeglm, variables = "beta[745]")
plot_coverage_diff(res3_apeglm, variables = "beta[745]")

empirical_coverage(dplyr::filter(res3_apeglm$stats, variable == "beta[745]"), c(0.5,0.9,0.95))

```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res3_apeglm, variables = pars_to_show)
plot_ecdf_diff(res3_apeglm, variables = pars_to_show)
plot_sim_estimated(res3_apeglm, variables = pars_to_show)
plot_coverage_diff(res3_apeglm, variables = pars_to_show)
```
```{r}
overall_coverage(res3_apeglm)
```



```{r}
coverage_plot <- plot_coverage(merge_stats(res3_apeglm),
       max_points = 101)
ggsave(file.path(image_path, "3_apeglm_coverage.png"), coverage_plot,
        width = image_width, height = image_height)
```

```{r}
ggsave(file.path(image_path,"3_apeglm_diff.png"),
  plot_coverage_diff_merged(res3_apeglm),
  width = image_width, height = image_height)
```



```{r}
# TODO resolve, why we have NA lfcSE with apeglm (maybe a bug in apeglm)
# tt <- SBC_fit.SBC_backend_DEseq2(SBC_backend_DESeq2(shrink_type = "apeglm"), datasets3$generated[[810]],  1)
# mm <- SBC_fit_to_draws_matrix.SBC_fit_DESeq2(tt)
# rr <- results(tt$dds)
# ss <- lfcShrink(tt$dds, coef = "group_B_vs_A", type = "apeglm")
# ss2 <- lfcShrink(tt$dds, coef = "group_B_vs_A", type = "normal")
# rr[613,]
# ss[613,]
# ss2[613,]

```

```{r}
datasets3$generated[[613]]$countData[613,]
```

## 3 samples, apeglm, use T

```{r}
res3_apeglm_t <- compute_SBC(datasets3, SBC_backend_DESeq2(shrink_type = "apeglm", useT = TRUE), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3_apeglm_t"),
                        keep_fits = FALSE)
```

```{r}
overall_coverage(res3_apeglm_t)
```


```{r}
p <- plot_coverage_diff_merged(res3_apeglm_t)
p
ggsave(file.path(image_path,"3_apeglm_t_diff.png"), p,
  width = image_width, height = image_height)
```


## 3 samples, narrow, none

```{r}
set.seed(585222)
datasets3_nar <- generate_datasets(SBC_generator_function(
   generator_single_DEseq2, N_per_group = 3, N_genes = 1000, predictor_prior_sd = 0.5),
   n_sims = 1000) 
```

## 3 samples narrow, none

```{r}
res3_nar_none <- compute_SBC(datasets3_nar, SBC_backend_DESeq2(shrink_type = "none"), globals = c("SBC_fit.SBC_backend_DEseq2", "SBC_fit_to_draws_matrix.SBC_fit_DESeq2",
                                                          "SBC_backend_iid_draws.SBC_backend_DEseq2"),
                         cache_mode = "results", cache_location = file.path(cache_dir, "SBC_3_nar_none"),
                        keep_fits = FALSE)

```


```{r}
pars_to_show <- paste0("beta[", sample(1:1000, size = 12), "]")
plot_rank_hist(res3_nar_none, variables = pars_to_show)
plot_ecdf_diff(res3_nar_none, variables = pars_to_show)
plot_coverage_diff(res3_nar_none, variables = pars_to_show)
```

```{r}
overall_coverage(res3_nar_none)
```