Simulation analysis: systematic and for all lines

knitr::opts_chunk$set(echo = TRUE, warning = FALSE, message = FALSE)
dir.create("figures/simulations", showWarnings = FALSE, recursive = TRUE)
library(ggpubr)
library(tidyverse)
library(cardelino)
library(viridis)
library(cowplot)
library(latex2exp)
library(ggrepel)
lines <- c("euts", "fawm", "feec", "fikt", "garx", "gesg", "heja", "hipn", 
            "ieki", "joxm", "kuco", "laey", "lexy", "naju", "nusw", "oaaz", 
            "oilg", "pipw", "puie", "qayj", "qolg", "qonc", "rozh", "sehl", 
            "ualf", "vass", "vuna", "wahn", "wetu", "xugn", "zoxy", "vils")

# assess cardelino with simulation
dat_dir <- "data/"
data(config_all)
data(simulation_input)
simu_input <- list("D" = D_input)

demuxlet <- function(A, D, Config, theta1 = 0.5, theta0 = 0.01) {
  P0_mat <- dbinom(A, D, theta0, log = TRUE)
  P1_mat <- dbinom(A, D, theta1, log = TRUE)
  P0_mat[which(is.na(P0_mat))] <- 0
  P1_mat[which(is.na(P1_mat))] <- 0
  logLik_mat <- t(P0_mat) %*% (1 - Config) + t(P1_mat) %*% Config
  prob_mat <- exp(logLik_mat) / rowSums(exp(logLik_mat))
  prob_mat
}

simulate_joint <- function(Config_all, D_all, n_clone = 4, mut_size = 5,
                           missing = NULL, error_mean = c(0.01, 0.44), 
                           error_var = c(30, 4.8), n_repeat = 1) {
  simu_data_full <- list()
  for (i in seq_len(n_repeat)) {
    Config <- sample(Config_all[[n_clone - 2]], size = 1)[[1]]
    Config <- matrix(rep(c(t(Config)), mut_size), ncol = ncol(Config), 
                     byrow = TRUE)
    row.names(Config) <- paste0("SNV", seq_len(nrow(Config)))
    colnames(Config) <- paste0("Clone", seq_len(ncol(Config)))
    D_input <- sample_seq_depth(D_all, n_cells = 200, missing_rate = missing, 
                                n_sites = nrow(Config))
    sim_dat <- sim_read_count(Config, D_input, Psi = NULL, cell_num = 200, 
                              means = error_mean, vars = error_var)
    sim_dat[["Config"]] <- Config
    simu_data_full[[i]] <- sim_dat
  }
  simu_data_full
}

assign_score <- function(prob_mat, simu_mat, threshold=0.2, mode="delta") {
  assign_0 <- cardelino::get_prob_label(simu_mat)
  assign_1 <- cardelino::get_prob_label(prob_mat)
  prob_val <- cardelino::get_prob_value(prob_mat, mode = mode)
  idx <- prob_val >= threshold
  
  rt_list <- list("ass" = mean(idx), 
                  "acc" = mean(assign_0 == assign_1),
                  "acc_ass" = mean((assign_0 == assign_1)[idx]))
  rt_list
}

assign_curve <- function(prob_mat, simu_mat, mode="delta"){
  assign_0 <- cardelino::get_prob_label(simu_mat)
  assign_1 <- cardelino::get_prob_label(prob_mat)
  prob_val <- cardelino::get_prob_value(prob_mat, mode = mode)
  
  thresholds <- sort(unique(prob_val))
  ACC <- rep(0, length(thresholds))
  ASS <- rep(0, length(thresholds))
  for (i in seq_len(length(thresholds))) {
    idx <- prob_val >= thresholds[i]
    ASS[i] <- mean(idx)
    ACC[i] <- mean((assign_0 == assign_1)[idx])
  }
  thresholds <- c(thresholds, 1.0)
  ACC <- c(ACC, 1.0)
  ASS <- c(ASS, 0.0)
  AUC <- AUC_acc_ass <- 0.0
  for (i in seq_len(length(thresholds) - 1)) {
    AUC <- AUC + 0.5 * (thresholds[i] - thresholds[i + 1]) * 
      (ACC[i] + ACC[i + 1])
    AUC_acc_ass <- AUC_acc_ass + 0.5 * (ASS[i] - ASS[i + 1]) * 
      (ACC[i] + ACC[i + 1])
  }
  AUC <- AUC / (thresholds[1] - thresholds[length(thresholds)])
  AUC_acc_ass <- AUC_acc_ass / (ASS[1] - ASS[length(thresholds)])
  
  rt_list <- list("ACC" = ACC, "ASS" = ASS, "AUC" = AUC, 
                  "AUC_acc_ass" = AUC_acc_ass, "thresholds" = thresholds)
  rt_list
}

assign_macro_ROC <- function(prob_mat, simu_mat) {
  thresholds <- seq(0, 0.999, 0.001)
  ACC <- rep(0, length(thresholds))
  ASS <- rep(0, length(thresholds))
  FPR <- rep(0, length(thresholds))
  TPR <- rep(0, length(thresholds))
  for (i in seq_len(length(thresholds))) {
    idx <- prob_mat >= thresholds[i]
    ASS[i] <- mean(idx) # not very meaningful
    ACC[i] <- mean(simu_mat[idx])
    FPR[i] <- sum(simu_mat[idx] == 0) / sum(simu_mat == 0)
    TPR[i] <- sum(simu_mat[idx] == 1) / sum(simu_mat == 1)
  }
  
  AUC <- 0.0
  for (i in seq_len(length(thresholds) - 1)) {
    AUC <- AUC + 0.5 * (FPR[i] - FPR[i + 1]) * (TPR[i] + TPR[i + 1])
  }
  
  rt_list <- list("FPR" = FPR, "TPR" = TPR, "AUC" = AUC,
                  "thresholds" = thresholds, "ACC" = ACC, "ASS" = ASS)
  rt_list
}

set.seed(1)
ACC_all <- c()
ASS_all <- c()
AUC_all <- c()
ERR_all <- c()
labels_all <- c()
method_all <- c()
variable_all <- c()

type_use <- c("mut_size", "n_clone", "missing", "FNR", "shapes1")
value_list <- list(c(3, 5, 7, 10, 15, 25), 
                   seq(3, 8), 
                   seq(0.7, 0.95, 0.05), 
                   seq(0.35, 0.6, 0.05),
                   c(0.5, 1.0, 2.0, 4.0, 8.0, 16.0))

for (mm in seq_len(length(value_list))) {
  values_all <- value_list[[mm]]
  for (k in seq_len(length(values_all)))  {
    if (mm == 1) {
      simu_data <- simulate_joint(Config_all, simu_input$D, n_clone = 4, 
                                  mut_size = values_all[k], missing = 0.8, 
                                  error_mean = c(0.01, 0.44), n_repeat = 20)
    } else if (mm == 2) { 
      simu_data <- simulate_joint(Config_all, simu_input$D, 
                                  n_clone = values_all[k], mut_size = 10, 
                                  missing = 0.8, error_mean = c(0.01, 0.44),
                                  n_repeat = 20)
    } else if (mm == 3) {
      simu_data <- simulate_joint(Config_all, simu_input$D, n_clone = 4, 
                                  mut_size = 10, missing = values_all[k], 
                                  error_mean = c(0.01, 0.44), n_repeat = 20)
    } else if (mm == 4) {
      simu_data <- simulate_joint(Config_all, simu_input$D, n_clone = 4, 
                                  mut_size = 10, missing = 0.8, 
                                  error_mean = c(0.01, values_all[k]), 
                                  n_repeat = 20)
    } else if (mm == 5) {
      simu_data <- simulate_joint(Config_all, simu_input$D, n_clone = 4, 
                                  mut_size = 10, missing = 0.8, 
                                  error_mean = c(0.01, 0.44), 
                                  error_var = c(30, values_all[k]), 
                                  n_repeat = 20)
    }
    
    for (d_tmp in simu_data) {
      prob_all <- list()
      methods_use <- c("demuxlet", "Bern_EM", "Binom_EM", "Binom_Gibbs",
                       "Binom_gmline")
      prob_all[[1]] <- demuxlet(d_tmp$A_sim, d_tmp$D_sim, d_tmp$Config, 
                                theta0 = mean(d_tmp$theta0_binom, na.rm = TRUE))
      prob_all[[2]] <- cell_assign_EM(d_tmp$A_sim, d_tmp$D_sim, d_tmp$Config, 
                                      Psi = NULL, verbose = F)$prob
      prob_all[[3]] <- cell_assign_EM(d_tmp$A_sim, d_tmp$D_sim, d_tmp$Config, 
                                      Psi = NULL, model = "binomial",
                                      verbose = FALSE)$prob
      prob_all[[4]] <- cell_assign_Gibbs(d_tmp$A_sim, d_tmp$D_sim,
                                         d_tmp$Config, Psi = NULL,
                                         min_iter = 1000, wise = "variant", 
                                         prior1 = c(2.11, 2.69), 
                                         verbose = FALSE)$prob
      
      for (i in seq_len(length(prob_all))) {
        prob_mat <- prob_all[[i]]
        assign_scr <- assign_score(prob_mat, d_tmp$I_sim, 
                                                threshold = 0.5001, 
                                                mode = "best")
        ACC_all <- c(ACC_all, assign_scr$acc_ass)
        ASS_all <- c(ASS_all, assign_scr$ass)
        AUC_all <- c(AUC_all, assign_curve(prob_mat, d_tmp$I_sim, 
                                           mode = "best")$AUC_acc_ass)
        ERR_all <- c(ERR_all, mean(abs(prob_mat - d_tmp$I_sim)))
        labels_all <- c(labels_all, values_all[k])
        method_all <- c(method_all, methods_use[i])
        variable_all <- c(variable_all, type_use[mm])
      }
    }
  }
}

df <- data.frame(Accuracy = ACC_all, AUC_of_ACC_ASS = AUC_all, 
                 Assignable = ASS_all, MAE = ERR_all, Methods = method_all, 
                 labels = labels_all, variable = variable_all)

dat_dir <- "data/simulations"
saveRDS(df, paste0(dat_dir, "/simulate_extra_s1_v2.rds"))

set.seed(1)
PCAU_curve_list <- list()

simu_data <- simulate_joint(Config_all, simu_input$D, n_clone = 4, 
                              mut_size = 10, missing = 0.8, 
                              error_mean = c(0.01, 0.44), n_repeat = 20)

assign_0 <- matrix(0, nrow = 200, ncol = 20)
assign_1 <- matrix(0, nrow = 200, ncol = 20)
prob_all <- matrix(0, nrow = 200, ncol = 20)
for (i in seq_len(length(simu_data))) {
  d_tmp <- simu_data[[i]]
  prob_tmp <- cell_assign_Gibbs(d_tmp$A_sim, d_tmp$D_sim,
                                d_tmp$Config, Psi = NULL,
                                min_iter = 1000, wise = "variant", 
                                prior1 = c(2.11, 2.69), verbose = FALSE)$prob
  assign_0[, i] <- get_prob_label(d_tmp$I_sim)
  assign_1[, i] <- get_prob_label(prob_tmp)
  prob_all[, i] <- get_prob_value(prob_tmp, mode = "best")
}

dat_dir = "data/simulations"
saveRDS(list("assign_0" = assign_0, "assign_1" = assign_1, 
             "prob_all" = prob_all), 
        paste0(dat_dir, "/simulate_prob_curve.rds"))

p <- df %>%
  dplyr::mutate(labs = ifelse(line == "joxm", "joxm", "")) %>%
  dplyr::mutate(sites_per_clone = cut(n_sites / pmax(n_clone - 1, 1), 
                                      breaks = c(0, 3, 8, 15, 25, 60))) %>%
  ggplot(
         aes(x = recall_simu, y = recall_real, 
             fill = sites_per_clone)) +
  geom_abline(slope = 1, intercept = 0, colour = "gray40", linetype = 2) +
  geom_smooth(aes(group = 1), method = "lm", colour = "firebrick") +
  geom_label_repel(aes(label = labs), fill = "gray90", size = 4,
                   nudge_x = 0.05, nudge_y = -0.1) +
  geom_point(size = 3, shape = 21) +
  xlim(0, 1) + ylim(0, 1) +
  scale_fill_manual(name = "mean\n# variants\nper clonal\nbranch", 
                    values = magma(6)[-1]) +
  guides(colour = FALSE, group = FALSE) +
  xlab("Assignment rate: simulated") +
  ylab("Assignment rate: observed") +
  theme_cowplot(font_size = 16)

ggsave("figures/simulations/assign_rate_obs_v_sim.png", plot = p,
       height = 4.5, width = 5)
ggsave("figures/simulations/assign_rate_obs_v_sim.pdf", plot = p,
       height = 4.5, width = 5)

ggsave("figures/simulations/assign_rate_obs_v_sim_wide.png", plot = p,
       height = 4.5, width = 6.5)
ggsave("figures/simulations/assign_rate_obs_v_sim_wide.pdf", plot = p,
       height = 4.5, width = 6.5)
p

ggsave("figures/simulations/assign_rate_obs_v_sim_skinny.png", 
       plot = p + theme(legend.position = "bottom"),
       height = 5.5, width = 5)
ggsave("figures/simulations/assign_rate_obs_v_sim_skinny.pdf", 
       plot = p + theme(legend.position = "bottom"),
       height = 5.5, width = 5)

summary(lm(recall_real ~ recall_simu, data = df))

df %>%
  dplyr::mutate(sites_per_clone = cut(n_sites / n_clone, 
                                      breaks = c(0, 5, 10, 20, 40))) %>%
  ggplot(
         aes(x = recall_simu, y = precision_simu, 
             fill = sites_per_clone)) +
  geom_hline(yintercept = 0.85, colour = "gray40", linetype = 2) +
  geom_smooth(aes(group = 1), method = "lm", colour = "firebrick") +
  geom_point(size = 3, shape = 21) +
  xlim(0, 1) + ylim(0, 1) +
  scale_fill_manual(name = "mean\n# variants\nper clone", 
                    values = magma(5)[-1]) +
  guides(colour = FALSE, group = FALSE) +
  xlab("Assignment rate (recall)") +
  ylab("Precision")

ggsave("figures/simulations/sim_precision_v_recall.png", 
       height = 4.5, width = 5.5)
ggsave("figures/simulations/sim_precision_v_recall.pdf", 
       height = 4.5, width = 5.5)

table(df$n_clone)

summary(df$n_sites / (df$n_clone - 1))

devtools::session_info()