R/mts_plotClusterDistribution.R

In MultiSEp: Predict Synthetic Lethality and Other Gene Dependency Relationships from Multiomics Data

globalVariables(c("Gene1", "Gene2", "Mode1",
                  "Mode2", "PointFill", "Combo", "element_line"))

mts_plotClusterDistribution <- function(mixModelClusters, mixModelClusters2 = NULL,
                                        TScluster1 = NULL, TScluster2 = NULL,
                                        gene1, gene2,
                                        suppressBlankCanvas = NULL,
                                        gene1_datatype = "", gene2_datatype = "") {

  ## ---- internal controls --------------------------------------------------
  x_transform      <- "auto"        
  y_transform      <- "auto"     
  transform_method <- "sqrt"        
  ratio_threshold  <- 12        
  skew_threshold   <- 0.20         
  pseudo_log_sigma <- 1             
  edge_pad         <- 0.015         
  xlim             <- NULL         
  ylim             <- NULL        
  point_alpha      <- 1             
  legend_position  <- "right"       
  legend_ncol      <- 1             
  legend_key_size  <- 10         
  legend_text_size <- 8             

  ## ---- input checks ----
  if (missing(mixModelClusters) || !is.list(mixModelClusters))
    stop("No mixModelClusters list provided: Please run mts_mixModelCluster to obtain mixModelClusters")
 
  if (missing(gene1)) stop("gene1 must be provided")
  if (missing(gene2)) stop("gene2 must be provided")
  
  if (!is.null(mixModelClusters2)) {
    if (!is.list(mixModelClusters2))
      stop("mixModelClusters2 must be a list object: Please run mts_mixModelCluster to obtain mixModelClusters2")
    if (!gene1 %in% names(mixModelClusters))
      stop(paste(gene1, " is not in supplied in the mixModelClusters list", sep = ""))
    if (!gene2 %in% names(mixModelClusters2))
      stop(paste(gene2, " is not in supplied in the mixModelClusters2 list", sep = ""))
    if (is.atomic(mixModelClusters[[gene1]]))
      stop(paste(gene1, " from the mixModelClusters list does not have cluster assignment, please choose another gene", sep = ""))
    if (is.atomic(mixModelClusters2[[gene2]]))
      stop(paste(gene2, " from the mixModelClusters2 list does not have cluster assignment, please choose another gene", sep = ""))
    tmp <- GMM_CCL(mixModelClusters, mixModelClusters2)
    mixModelClusters  <- tmp[[1]]
    mixModelClusters2 <- tmp[[2]]
    tab1 <- mixModelClusters[[gene1]]
    tab2 <- mixModelClusters2[[gene2]]
  } else {
    if (!length(which(names(mixModelClusters) %in% gene1)) > 0)
      stop(paste(gene1, " is not in supplied in the mixModelClusters list", sep = ""))
    if (!length(which(names(mixModelClusters) %in% gene2)) > 0)
      stop(paste(gene2, " is not in supplied in the mixModelClusters list", sep = ""))
    if (is.atomic(mixModelClusters[[gene1]]))
      stop(paste(gene1, " from the mixModelClusters list does not have cluster assignment, please choose another gene", sep = ""))
    if (is.atomic(mixModelClusters[[gene2]]))
      stop(paste(gene2, " from the mixModelClusters list does not have cluster assignment, please choose another gene", sep = ""))
    mixModelClusters2 <- mixModelClusters
    tab1 <- mixModelClusters[[gene1]]
    tab2 <- mixModelClusters2[[gene2]]
  }
  colnames(tab1) <- colnames(tab2) <- c("Sample", "log2_expression", "mode")

  merged_data <- merge(tab1, tab2, by = "Sample", suffixes = c("_gene1", "_gene2"))
  data <- data.frame(
    Sample = merged_data$Sample,
    Gene1  = as.numeric(merged_data$log2_expression_gene1),
    Gene2  = as.numeric(merged_data$log2_expression_gene2),
    Mode1  = as.character(merged_data$mode_gene1),
    Mode2  = as.character(merged_data$mode_gene2)
  )

  ## ---- palette ----
  gg_color_hue <- function(n) hcl(h = seq(15, 375, length = n + 1), l = 65, c = 100)[1:n]
  modes <- sort(unique(c(data$Mode1, data$Mode2)))
  pal   <- gg_color_hue(length(modes)); names(pal) <- modes

  ## ---- cluster cut points (boundary lines AND breaks) ----
  x_cuts <- if (length(unique(tab1$mode)) > 1)
    sort(sapply(seq_len(length(unique(tab1$mode)) - 1),
                function(i) min(data$Gene1[data$Mode1 == as.character(i + 1)]))) else numeric(0)
  y_cuts <- if (length(unique(tab2$mode)) > 1)
    sort(sapply(seq_len(length(unique(tab2$mode)) - 1),
                function(i) min(data$Gene2[data$Mode2 == as.character(i + 1)]))) else numeric(0)

  ## ---- resolve transforms per axis (gene1/gene2-agnostic) ----
  if (x_transform == "auto")
    x_transform <- as.character(mts_pickAxisTransform(data$Gene1, cuts = x_cuts, method = transform_method,
                                                      ratio_threshold = ratio_threshold, skew_threshold = skew_threshold))
  if (y_transform == "auto")
    y_transform <- as.character(mts_pickAxisTransform(data$Gene2, cuts = y_cuts, method = transform_method,
                                                      ratio_threshold = ratio_threshold, skew_threshold = skew_threshold))

  ## ---- limits, breaks, expansion ----
  xlims <- .mts_axis_limits(data$Gene1, data$Mode1, x_transform, override = xlim)
  ylims <- .mts_axis_limits(data$Gene2, data$Mode2, y_transform, override = ylim)

  x_breaks <- round(x_cuts, 1); x_breaks <- x_breaks[x_breaks >= xlims[1] & x_breaks <= xlims[2]]
  y_breaks <- round(y_cuts, 1); y_breaks <- y_breaks[y_breaks >= ylims[1] & y_breaks <= ylims[2]]
  if (length(x_breaks) == 0) x_breaks <- ggplot2::waiver()
  if (length(y_breaks) == 0) y_breaks <- ggplot2::waiver()

  ## Limits already reach the density tails, so expand flush (0) on the low end so
  ## the gaussian tail starts on the axis. sqrt floors at 0, so it keeps a small
  ## edge_pad to inset the lowest marker; tiny high pad so the top tail isn't clipped.
  x_low <- if (identical(x_transform, "sqrt")) edge_pad else 0
  y_low <- if (identical(y_transform, "sqrt")) edge_pad else 0
  x_expand <- ggplot2::expansion(mult = c(x_low, 0.01))
  y_expand <- ggplot2::expansion(mult = c(y_low, 0.01))

  ## ---- marginal density panels (same transform as their scatter axis) ----
  xaxis <- ggplot(data, aes(x = Gene1, fill = Mode1)) +
    geom_density(alpha = .3) +
    scale_fill_manual(values = pal, guide = "none") +
    .mts_axis_scale("x", x_transform, limits = xlims, breaks = NULL, expand = x_expand,
                    pseudo_log_sigma = pseudo_log_sigma) +
    ggplot2::scale_y_continuous(breaks = NULL) +
    labs(x = NULL, y = NULL) + theme_minimal() +
    theme(plot.margin = unit(c(0, 0, -35, 0), "pt"), panel.grid = element_blank())

  ## Gene2 is on the x aesthetic here (coord_flip-ped), so y_transform is applied
  ## through the x scale of this sub-panel.
  yaxis <- ggplot(data, aes(x = Gene2, fill = Mode2)) +
    geom_density(alpha = .3) +
    scale_fill_manual(values = pal, guide = "none") +
    .mts_axis_scale("x", y_transform, limits = ylims, breaks = NULL, expand = y_expand,
                    pseudo_log_sigma = pseudo_log_sigma) +
    coord_flip() + scale_y_reverse(breaks = NULL) +
    labs(x = NULL, y = NULL) + theme_minimal() +
    theme(plot.margin = unit(c(0, -35, 0, 0), "pt"), panel.grid = element_blank())

  ## ---- sample lists for tissue-specific colouring ----
  pull_samples <- function(lst) {
    if (is.null(lst)) return(character(0))
    unlist(lapply(lst, `[`, 1), use.names = FALSE)
  }
  ts1 <- unique(pull_samples(TScluster1))
  ts2 <- unique(pull_samples(TScluster2))

  ## ---- scatter base (TS branch or all-sample branch) ----
  if (!is.null(TScluster1) || !is.null(TScluster2)) {
    data$PointFill <- ifelse(data$Sample %in% ts1, "TS1",
                             ifelse(data$Sample %in% ts2, "TS2", "Other"))
    scatter <- ggplot() +
      geom_point(data = subset(data, PointFill == "Other"), aes(x = Gene1, y = Gene2),
                 shape = 21, size = 2, fill = "white", color = "black", alpha = point_alpha) +
      geom_point(data = subset(data, PointFill == "TS2"), aes(x = Gene1, y = Gene2),
                 shape = 21, size = 2, fill = "#0000FF", color = "black", alpha = point_alpha) +
      geom_point(data = subset(data, PointFill == "TS1"), aes(x = Gene1, y = Gene2),
                 shape = 21, size = 2, fill = "#FF0000", color = "black", alpha = point_alpha) +
      labs(x = paste(gene1, gene1_datatype), y = paste(gene2, gene2_datatype)) +
      theme(axis.text = element_text(size = 9), axis.title = element_text(size = 12),
            panel.background = element_blank(), panel.grid = element_blank(),
            axis.line = element_line(color = "black"), axis.ticks = element_line(color = "black"))
  } else {
    data$Combo <- with(data, factor(paste0("X", Mode1, ",Y", Mode2),
                                     levels = unique(paste0("X", Mode1, ",Y", Mode2))))
    combo_cols <- gg_color_hue(length(levels(data$Combo)))
    names(combo_cols) <- levels(data$Combo)
    scatter <- ggplot(data, aes(x = Gene1, y = Gene2, fill = Combo)) +
      geom_point(shape = 21, size = 2, color = "black", alpha = point_alpha) +
      scale_fill_manual(name = "Cluster combination", values = combo_cols,
                        guide = guide_legend(ncol = legend_ncol)) +
      labs(x = paste(gene1, gene1_datatype), y = paste(gene2, gene2_datatype)) +
      theme(axis.text = element_text(size = 9), axis.title = element_text(size = 12),
            panel.background = element_blank(),
            panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
            axis.line = element_line(color = "black"), axis.ticks = element_line(color = "black"),
            legend.position = "right")
  }

  ## scatter shows regular numeric ticks PLUS the cut (abline) labels; the marginals
  ## keep just the cut labels (x_breaks / y_breaks) 
  n_mode1 <- length(unique(tab1$mode))
  n_mode2 <- length(unique(tab2$mode))
  xb_scatter <- if (n_mode1 > 2) x_breaks else .mts_combine_breaks(x_cuts, xlims, x_transform)
  yb_scatter <- if (n_mode2 > 2) y_breaks else .mts_combine_breaks(y_cuts, ylims, y_transform)

  ## ---- apply the per-axis transforms to the scatter ----
  ## clip = "off": low-end expansion is 0 (gaussians flush on the axis), which would
  ## otherwise bisect a marker on the min boundary; turning off clipping renders every
  ## point in full while the shared scale keeps the marginals aligned with the scatter.
  scatter <- scatter +
    .mts_axis_scale("x", x_transform, limits = xlims, breaks = xb_scatter, expand = x_expand,
                    pseudo_log_sigma = pseudo_log_sigma) +
    .mts_axis_scale("y", y_transform, limits = ylims, breaks = yb_scatter, expand = y_expand,
                    pseudo_log_sigma = pseudo_log_sigma) +
    coord_cartesian(clip = "off")

  ## ---- cluster boundary lines (both branches) ----
  for (cx in x_cuts) scatter <- scatter + geom_vline(xintercept = cx, color = "grey57", linewidth = 0.5)
  for (cy in y_cuts) scatter <- scatter + geom_hline(yintercept = cy, color = "grey57", linewidth = 0.5)

  ## ---- assemble ----
  ## Single 2x2 design (A = top density, B = left density, C = scatter, # = empty
  ## corner). A single patchwork assembly aligns its panels on its own; the nested
  ## (a+b)/(c+d) form let the all-sample branch's fill legend shift the top density
  ## off the scatter. axes = "keep" (NOT "collect") preserves the scatter's own
  ## cluster-cut break labels (collect would treat them as duplicates and drop them).
  ## `&` applies the legend settings to the collected guide.
  combined_plot <- xaxis + yaxis + scatter +
    plot_layout(design = "#A\nBC", widths = c(0.2, 1), heights = c(0.3, 1),
                guides = "collect", axes = "keep") &
    theme(legend.position = legend_position,
          legend.key.size = unit(legend_key_size, "pt"),
          legend.text     = element_text(size = legend_text_size),
          legend.title    = element_text(size = legend_text_size + 1),
          legend.spacing  = unit(2, "pt"))

  if (is.null(suppressBlankCanvas)) {
    return(combined_plot)
  } else {
    print(combined_plot)
  }
}