Predicting COO Using Mutation Data

mut_fisher <- 
  mut_features %>%
  dplyr::filter(coo_ns %in% c("ABC", "GCB")) %>%
  dplyr::group_by(feature) %>%
  tidyr::nest() %>%
  dplyr::mutate(fisher = purrr::map(data, ~fisher.test(.$mutated, .$coo_ns)),
                p_val = purrr::map_dbl(fisher, "p.value"),
                q_val = p.adjust(p_val, method = p_adjust_method)) %>%
  dplyr::select(-data, -fisher) %>%
  dplyr::arrange(q_val) %>%
  dplyr::filter(q_val < qval_cutoff)

knitr::kable(mut_fisher)

heatmap_mut_fisher <- 
  mut_feat_matrix %>% 
  extract(mut_fisher$feature, ) %>%
  pheatmap::pheatmap(cluster_rows = FALSE, cluster_cols = FALSE, show_colnames = FALSE,
                     color = colours$mutated, legend = FALSE,
                     annotation_col = heatmap_mut_annot,
                     annotation_colors = colours,
                     gaps_col = sum(heatmap_mut_annot$coo_ns == "ABC"))

# Remove any features with near-zero variance
nz_feats <- caret::nearZeroVar(mut_feat_df)
abc_gcb_rows <- mut_feat_df$coo_ns %in% c("ABC", "GCB")
mut_rf_patients <- mut_feat_df[abc_gcb_rows, 1]
mut_feat_df_no_nz <- mut_feat_df[abc_gcb_rows, -c(1, nz_feats)]

mut_ctrl <- caret::trainControl(method = "cv", number = 10)

mut_model <- 
  caret::train(coo_ns ~ ., mut_feat_df_no_nz, method = "rf",
               preProcess = "scale", trControl = mut_ctrl, tuneLength = 10)
mut_model

mut_imp <- caret::varImp(mut_model, scale = FALSE)

mut_rf <- 
  mut_imp$importance %>%
  tibble::rownames_to_column("feature") %>%
  dplyr::rename(importance = Overall) %>%
  dplyr::arrange(desc(importance)) %>%
  head(nrow(mut_fisher))

knitr::kable(mut_rf)

heatmap_mut_rf <- 
  mut_feat_matrix %>% 
  extract(mut_rf$feature, ) %>%
  pheatmap::pheatmap(cluster_rows = FALSE, cluster_cols = FALSE, show_colnames = FALSE,
                     color = colours$mutated, legend = FALSE,
                     annotation_col = heatmap_mut_annot,
                     annotation_colors = colours,
                     gaps_col = sum(heatmap_mut_annot$coo_ns == "ABC"))

gcb_features <- mut_hybrid_df_top %>% dplyr::filter(ratio < 1) %$% feature
abc_features <- mut_hybrid_df_top %>% dplyr::filter(ratio > 1) %$% feature

discretize <- function(x) {
  dplyr::recode(x, `0` = "0", `1` = "1", `2` = "2", .default = "3+")
}

num_feats <- 
  mut_features %>%
  dplyr::filter(mutated) %>%
  dplyr::group_by(patient, coo_ns) %>%
  dplyr::summarise(gcb_num_feats = discretize(sum(feature %in% gcb_features)),
                   abc_num_feats = discretize(sum(feature %in% abc_features))) 

heatmap_mut_hybrid_num_feats_annot <-
  num_feats %>%
  as.data.frame() %>%
  tibble::column_to_rownames("patient") %>%
  dplyr::ungroup()

barplot_mut_hybrid_num_feats <-
  heatmap_mut_hybrid_num_feats_annot %>%
  dplyr::select(-coo_ns,
                `Number of ABC features` = abc_num_feats,
                `Number of GCB features` = gcb_num_feats) %>%
  tidyr::gather(type, count) %>%
  dplyr::count(type, count) %>%
  ggplot(aes(x = count, y = n)) +
  geom_bar(stat = "identity") +
  facet_grid(~type) + 
  labs(title = "Number of patients per mutation feature count",
       x = "Number of ABC/GCB features", y = "Number of patients")
barplot_mut_hybrid_num_feats

# discretize_probs <- function(x) {
#   cut(x, breaks = c(0, 0.2, 0.4, 0.6, 0.8, 1), 
#       labels = c("0-20%", "20-40%", "40-60%", "60-80%", "80-100%")) %>%
#     as.character()
# }

discretize_diff <- function(x) {
  cut(x, breaks = c(-1.1, -0.6, -0.2, 0.2, 0.6, 1.1), 
      labels = c("GCB 60-100%", "GCB 20-60%", "Ambiguous", "ABC 20-60%", "ABC 60-100%")) %>%
    as.character()
}

# mut_rf_prob_df <- 
#   predict(mut_model, mut_feat_df_no_nz, "prob") %>%
#   dplyr::transmute(abc_rf_prob = discretize_probs(ABC), 
#                    gcb_rf_prob = discretize_probs(GCB)) %>%
#   dplyr::bind_cols(mut_rf_patients)

mut_rf_prob_df <- 
  predict(mut_model, mut_feat_df, "prob") %>%
  dplyr::transmute(abc_gcb_diff = discretize_diff(ABC - GCB)) %>%
  dplyr::bind_cols(mut_feat_df[,"patient"])

heatmap_mut_rf_prob_annot <-
  heatmap_mut_hybrid_num_feats_annot %>%
  tibble::rownames_to_column("patient") %>%
  dplyr::left_join(mut_rf_prob_df) %>%
  tibble::column_to_rownames("patient") %>%
  extract(, c("coo_ns", "gcb_num_feats", "abc_num_feats", "abc_gcb_diff"))

barplot_mut_rf_prob_data <-
  heatmap_mut_rf_prob_annot %>%
  dplyr::select(-dplyr::ends_with("_num_feats"),
                `P(ABC) - P(GCB)` = abc_gcb_diff) %>%
  tidyr::gather(type, prob, -coo_ns) %>%
  dplyr::count(coo_ns, type, prob) %>%
  dplyr::mutate(prob = factor(prob, names(colours$abc_gcb_diff)[1:5])) 

barplot_mut_rf_prob <-
  barplot_mut_rf_prob_data %>%
  ggplot(aes(x = prob, y = n, fill = coo_ns)) +
  geom_bar(stat = "identity") +
  scale_x_discrete(drop = FALSE) +
  scale_fill_manual(values = colours$coo_ns) +
  place_legend(c(1, 1)) +
  labs(title = "Number of patients per RF probability range",
       x = "P(ABC) - P(GCB)", y = "Number of patients", fill = "Lymph2Cx COO")
barplot_mut_rf_prob

patient_no_mut <-
  mut_feat_df %>% 
  dplyr::transmute(patient, total = rowSums(.[, -(1:2)])) %>% 
  dplyr::filter(total < 5) %$% 
  patient

heatmap_mut_rf_prob_annot_no_mut <-
  heatmap_mut_rf_prob_annot %>%
  tibble::rownames_to_column("patient") %>% 
  dplyr::filter(patient %in% patient_no_mut) %>%
  dplyr::select(-dplyr::ends_with("_num_feats"),
                `P(ABC) - P(GCB)` = abc_gcb_diff,
                -patient) %>%
  tidyr::gather(type, prob, -coo_ns) %>%
  dplyr::count(coo_ns, type, prob) %>%
  dplyr::mutate(prob = factor(prob, names(colours$abc_gcb_diff)[1:5]))

barplot_mut_rf_prob_no_mut <- 
  barplot_mut_rf_prob %+% 
  heatmap_mut_rf_prob_annot_no_mut +
  labs(title = paste0(barplot_mut_rf_prob$labels$title, "\n(Patients with <5 mutations)")) +
  place_legend(c(0.5, 1))
barplot_mut_rf_prob_no_mut

lineplot_mut_rf_prob_df <-
  predict(mut_model, mut_feat_df, "prob") %>%
  dplyr::bind_cols(mut_feat_df[,c("patient", "coo_ns")]) %>%
  tibble::rownames_to_column("id") %>%
  dplyr::ungroup() %>%
  dplyr::mutate(id = forcats::fct_reorder(id, ABC),
                prob = GCB)

scatterplot_rf_prob_vs_num_muts_df <- 
  mut_feat_df %>% 
  tidyr::gather(feature, num_muts, -patient, -coo_ns) %>% 
  dplyr::group_by(patient) %>% 
  dplyr::summarise(total_muts = sum(num_muts)) %>% 
  dplyr::left_join(lineplot_mut_rf_prob_df, by = "patient")

scatterplot_rf_prob_vs_num_muts <- 
  scatterplot_rf_prob_vs_num_muts_df %>% {
    ymax <- max(.$total_muts) + 2
    ggplot(., aes(x = GCB, y = total_muts, colour = coo_ns)) + 
    annotate("rect", xmin = 0, xmax = 0.5, ymin = 0, ymax = ymax,
             fill = colours$coo["ABC"], alpha = 0.1) +
    annotate("rect", xmin = 0.5, xmax = 1, ymin = 0, ymax = ymax, 
             fill = colours$coo["GCB"], alpha = 0.1) +
    geom_vline(xintercept = 0.5) +
    geom_point() +
    scale_x_continuous(limits = c(0, 1), expand = c(0, 0)) +
    scale_y_continuous(limits = c(0, ymax), expand = c(0, 0)) +
    scale_color_manual(values = colours$coo) +
    labs(title = paste0("Relationship between probability for GCB and number of mutations\n",
                        "(Background colour corresponds to random forest prediction)"), 
           x = "Probability for GCB", y = "Number of mutations", colour = "Lymph2Cx COO") +
    place_legend(c(0, 1))
  }
scatterplot_rf_prob_vs_num_muts

n_patients <- ncol(mut_feat_matrix_all)

lineplot_mut_rf_prob <- 
  lineplot_mut_rf_prob_df %>% 
  ggplot(aes(x = id, y = prob, colour = coo)) +
    annotate("rect", xmin = 0, xmax = n_patients, ymin = 0, ymax = 0.5, 
             fill = colours$coo["ABC"], alpha = 0.1) +
    annotate("rect", xmin = 0, xmax = n_patients, ymin = 0.5, ymax = 1, 
             fill = colours$coo["GCB"], alpha = 0.1) +
    geom_hline(yintercept = 0.5) +
    geom_hline(yintercept = c(0.35, 0.65), linetype = 3) +
    geom_point(aes(colour = coo_ns), size = 0.5) +
    scale_color_manual(values = colours$coo) +
    place_legend(c(0, 0)) +
    scale_x_discrete(labels = NULL) +
    scale_y_continuous(limits = c(0, 1), expand = c(0, 0)) +
    theme(axis.ticks.x = element_blank()) +
    labs(title = paste0("Random forest probabilities for GCB\n",
                        "(Background colour corresponds to random forest prediction)"), 
         x = "Patient", y = "Probability for GCB", colour = "Lymph2Cx COO")
lineplot_mut_rf_prob