260423_CoexpressionPlot.R

library(Seurat)
library(dplyr)
library(ggplot2)

CoexpressionPlot <- function(
    object,
    gene1,
    gene2,
    reduction = "umap",
    assay = NULL,
    layer = "data",
    group.by = NULL,
    expr.cutoff1 = 0,
    expr.cutoff2 = 0,
    pt.size = 1,
    alpha = 0.9,
    cols = c(
      "No expression" = "lightgrey",
      "Gene1 only" = "#3B78E7",
      "Gene2 only" = "#32CD32",
      "All co-expressed" = "red"
    ),
    label = TRUE,
    label.size = 4,
    label.color = "black",
    label.box = FALSE,
    title = NULL
) {
  if (!inherits(object, "Seurat")) {
    stop("object must be a Seurat object.")
  }
  
  if (is.null(assay)) {
    assay <- DefaultAssay(object)
  }
  
  if (!reduction %in% Reductions(object)) {
    stop(paste0("Reduction '", reduction, "' not found in object."))
  }
  
  expr_mat <- GetAssayData(object = object, assay = assay, layer = layer)
  
  missing_genes <- setdiff(c(gene1, gene2), rownames(expr_mat))
  if (length(missing_genes) > 0) {
    stop(
      paste0(
        "These gene(s) were not found in assay='", assay,
        "', layer='", layer, "': ",
        paste(missing_genes, collapse = ", ")
      )
    )
  }
  
  cells_use <- colnames(object)
  
  expr_df <- data.frame(
    cell = cells_use,
    gene1_expr = as.numeric(expr_mat[gene1, cells_use]),
    gene2_expr = as.numeric(expr_mat[gene2, cells_use]),
    stringsAsFactors = FALSE
  )
  
  emb <- Embeddings(object[[reduction]])
  emb <- as.data.frame(emb)
  
  if (ncol(emb) < 2) {
    stop(paste0("Reduction '", reduction, "' has fewer than 2 dimensions."))
  }
  
  emb <- emb[cells_use, 1:2, drop = FALSE]
  colnames(emb)[1:2] <- c("dim1", "dim2")
  emb$cell <- rownames(emb)
  
  df <- left_join(expr_df, emb, by = "cell")
  
  if (!is.null(group.by)) {
    if (!group.by %in% colnames(object@meta.data)) {
      stop(paste0("group.by column '", group.by, "' not found in object@meta.data"))
    }
    df$group_label <- object@meta.data[df$cell, group.by]
  }
  
  df <- df %>%
    mutate(
      expr_group = case_when(
        gene1_expr <= expr.cutoff1 & gene2_expr <= expr.cutoff2 ~ "No expression",
        gene1_expr >  expr.cutoff1 & gene2_expr <= expr.cutoff2 ~ "Gene1 only",
        gene1_expr <= expr.cutoff1 & gene2_expr >  expr.cutoff2 ~ "Gene2 only",
        gene1_expr >  expr.cutoff1 & gene2_expr >  expr.cutoff2 ~ "All co-expressed"
      )
    )
  
  df$expr_group <- factor(
    df$expr_group,
    levels = c("No expression", "Gene1 only", "Gene2 only", "All co-expressed")
  )
  
  df <- df %>% arrange(expr_group)
  
  legend_labels <- c(
    "No expression" = "No expression",
    "Gene1 only" = paste0(gene1, " only"),
    "Gene2 only" = paste0(gene2, " only"),
    "All co-expressed" = "All co-expressed"
  )
  
  if (is.null(title)) {
    title <- paste0(gene1, " & ", gene2, " co-expression")
  }
  
  p <- ggplot(df, aes(x = dim1, y = dim2, color = expr_group)) +
    geom_point(size = pt.size, alpha = alpha) +
    scale_color_manual(
      values = cols,
      breaks = names(legend_labels),
      labels = legend_labels
    ) +
    labs(
      title = title,
      x = paste0(toupper(reduction), " 1"),
      y = paste0(toupper(reduction), " 2"),
      color = NULL
    ) +
    theme_classic(base_size = 14)
  
  if (label && !is.null(group.by)) {
    label_df <- df %>%
      dplyr::filter(!is.na(group_label)) %>%
      dplyr::group_by(group_label) %>%
      dplyr::summarise(
        dim1 = median(dim1, na.rm = TRUE),
        dim2 = median(dim2, na.rm = TRUE),
        .groups = "drop"
      ) %>%
      dplyr::mutate(label = as.character(group_label))
    
    if (label.box) {
      p <- p +
        geom_label(
          data = label_df,
          mapping = aes(x = dim1, y = dim2, label = label),
          inherit.aes = FALSE,
          size = label.size,
          color = label.color
        )
    } else {
      p <- p +
        geom_text(
          data = label_df,
          mapping = aes(x = dim1, y = dim2, label = label),
          inherit.aes = FALSE,
          size = label.size,
          color = label.color,
          fontface = "bold"
        )
    }
  }
  
  return(p)
}