R/dotplot.R

In GuangchuangYu/enrichplot: Visualization of Functional Enrichment Result

##' @rdname dotplot
##' @exportMethod dotplot
##' @author Guangchuang Yu
setMethod("dotplot", signature(object = "enrichResult"),
          function(object, x = "GeneRatio", color = "p.adjust",
                   showCategory=10, size = NULL,
                   split = NULL, font.size=12, title = "",
                   orderBy="x", label_format = 30, ...) {
              dotplot.enrichResult(object = object, x = x, color = color,
                                   showCategory = showCategory,
                                   size = size, split = split,
                                   font.size = font.size,
                                   title = title, orderBy = orderBy,
                                   label_format = label_format, ...)
          })

##' @rdname dotplot
##' @exportMethod dotplot
setMethod("dotplot", signature(object = "gseaResult"),
          function(object, x = "GeneRatio", color = "p.adjust", showCategory=10,
                   size = NULL, split = NULL, font.size=12, title = "",
                   orderBy="x", label_format = 30, ...) {
                if (color == "NES") {
                    NES <- TRUE
                    color <- "p.adjust"
                } else {
                    NES <- FALSE
                }
                p <- dotplot.enrichResult(object = object, x = x, color = color,
                        showCategory = showCategory,
                        size = size, split = split,
                        font.size = font.size,
                        title = title, orderBy = orderBy,
                        label_format = label_format, ...)
                
                if (NES) {
                    p <- suppressMessages(p + aes_(fill=~NES) + 
                        # scale_fill_continuous(name = "NES") +
                        set_enrichplot_color(type = "fill", name = "NES")
                    )
                }
                return(p)
          })

##' @rdname dotplot
##' @aliases dotplot,compareClusterResult,ANY-method
##' @exportMethod dotplot
setMethod("dotplot", signature(object="compareClusterResult"),
          function(object,
                   x = "Cluster",
                   color ="p.adjust",
                   showCategory=5,
                   split=NULL,
                   font.size=12,
                   title="",
                   by="geneRatio",
                   size=NULL,
                   includeAll=TRUE,
                   label_format = 30,
                   ...
                   ) {
              dotplot.compareClusterResult(object, x=x, colorBy = color,
                                           showCategory = showCategory, by = by,
                                           size = size, includeAll = includeAll,
                                           split = split, font.size = font.size,
                                           title = title, label_format = label_format,
                                           ...)
})

##' @rdname dotplot
##' @exportMethod dotplot
##' @aliases dotplot,enrichResultList,ANY-method
##' @author Guangchuang Yu
setMethod("dotplot", signature(object = "enrichResultList"),
          function(object, x = "GeneRatio", color = "p.adjust",
                   showCategory=10, size = NULL,
                   split = NULL, font.size=12, title = "",
                   orderBy="x", label_format = 30, ...) {
              dotplot.enrichResult(object = object, x = x, color = color,
                                   showCategory = showCategory,
                                   size = size, split = split,
                                   font.size = font.size,
                                   title = title, orderBy = orderBy,
                                   label_format = label_format, ...)
})

##' @rdname dotplot
##' @exportMethod dotplot
##' @aliases dotplot,gseaResultList,ANY-method
setMethod("dotplot", signature(object = "gseaResultList"),
          function(object, x = "GeneRatio", color = "p.adjust",
                   showCategory=10, size = NULL,
                   split = NULL, font.size=12, title = "",
                   orderBy="x", label_format = 30, ...) {
              if (color == "NES") {
                      NES <- TRUE
                      color <- "p.adjust"
              } else {
                      NES <- FALSE
              }
              p <- dotplot.enrichResult(object = object, x = x, color = color,
                      showCategory = showCategory,
                      size = size, split = split,
                      font.size = font.size,
                      title = title, orderBy = orderBy,
                      label_format = label_format, ...)
              
              if (NES) {
                  p <- suppressMessages(p + aes_(fill=~NES) + 
                      # scale_fill_continuous(name = "NES") + 
                      set_enrichplot_color(type = "fill", name = "NES")
                  )
              }
              return(p)
})

#' @rdname dotplot
#' @param x variable for x-axis, one of 'GeneRatio' and 'Count'
#' @param color variable that used to color enriched terms,
#'              e.g. 'pvalue', 'p.adjust' or 'qvalue'
#' @param showCategory A number or a list of terms. If it is a number,
#' the first n terms will be displayed. If it is a list of terms,
#' the selected terms will be displayed.
#' @param size variable that used to scale the sizes of categories,
#' one of "geneRatio", "Percentage" and "count"
#' @param split separate result by 'category' variable
#' @param font.size font size
#' @param title plot title
#' @param label_format a numeric value sets wrap length, alternatively a
#' custom function to format axis labels.
#' by default wraps names longer that 30 characters
#' @param orderBy The order of the Y-axis
#' @param decreasing logical. Should the orderBy order be increasing or decreasing?
#' @importFrom ggplot2 fortify
#' @importFrom ggplot2 ggplot
#' @importFrom ggplot2 aes_string
#' @importFrom ggplot2 geom_point
#' @importFrom ggplot2 scale_color_gradient
#' @importFrom ggplot2 scale_color_continuous
#' @importFrom ggplot2 xlab
#' @importFrom ggplot2 ylab
#' @importFrom ggplot2 labs
#' @importFrom ggplot2 ggtitle
#' @importFrom ggplot2 scale_y_discrete
#' @importFrom ggplot2 guides
#' @importFrom ggplot2 guide_legend
#' @importFrom methods is
dotplot.enrichResult <- function(object, x = "geneRatio", color = "p.adjust",
                             showCategory=10, size=NULL, split = NULL,
                             font.size=12, title = "", orderBy="x",
                             label_format = 30, decreasing=TRUE) {
    colorBy <- match.arg(color, c("pvalue", "p.adjust", "qvalue"))
    if (x == "geneRatio" || x == "GeneRatio") {
        x <- "GeneRatio"
        if (is.null(size))
            size <- "Count"
    } else if (x == "count" || x == "Count") {
        x <- "Count"
        if (is.null(size))
            size <- "GeneRatio"
    } else if (is(x, "formula")) {
        x <- as.character(x)[2]
        if (is.null(size))
            size <- "Count"
    } else {
        ## message("invalid x, setting to 'GeneRatio' by default")
        ## x <- "GeneRatio"
        ## size <- "Count"
        if (is.null(size))
            size  <- "Count"
    }
    
    if (inherits(object, c("enrichResultList", "gseaResultList"))) {
        ldf <- lapply(object, fortify, showCategory=showCategory, split=split)
        df <- dplyr::bind_rows(ldf, .id="category")
        df$category <- factor(df$category, levels=names(object))
    } else {
        df <- fortify(object, showCategory = showCategory, split=split)
        ## already parsed in fortify
        ## df$GeneRatio <- parse_ratio(df$GeneRatio)
    }

    if (orderBy !=  'x' && !orderBy %in% colnames(df)) {
        message('wrong orderBy parameter; set to default `orderBy = "x"`')
        orderBy <- "x"
    }

    if (orderBy == "x") {
        df <- dplyr::mutate(df, x = eval(parse(text=x)))
    }

    label_func <- .label_format(label_format)

    idx <- order(df[[orderBy]], decreasing = decreasing)

    df$Description <- factor(df$Description,
                          levels=rev(unique(df$Description[idx])))
    

    p <- ggplot(df, aes(x=.data[[x]], y=.data[["Description"]], 
                        size=.data[[size]], fill=.data[[colorBy]])) +
        geom_point() +
        aes(shape = I(enrichplot_point_shape)) + 
        # scale_fill_continuous(name = color) +
        set_enrichplot_color(type = "fill", name = color) + 
        scale_y_discrete(labels = label_func) +
        ylab(NULL) + ggtitle(title) + theme_dose(font.size) 
        
    if (size == "Count") {
        # integer
        #size_n <- ceiling(sqrt(length(unique(df[[size]]))))
        size_break <- pretty(df[[size]], n=4)
        p <- p + scale_size(range=c(3, 8), breaks = size_break)
    } else {
        p <- p + scale_size(range=c(3, 8))
    }
    
    class(p) <- c("enrichplotDot", class(p))
    return(p)
}

.label_format <- function(label_format = 30) {
    label_func <- default_labeller(label_format)
    if(is.function(label_format)) {
        label_func <- label_format
    } 
    return(label_func)   
}


#' @rdname dotplot
#' @param object compareClusterResult object
#' @param by one of "geneRatio", "Percentage" and "count"
#' @param split apply `showCategory` to each category specified by the 'split', e.g., "ONTOLOGY", "category" and "intersect".  Default is NULL and do nothing
#' @param includeAll logical
#' @param font.size font size
#' @param title figure title
#' @param group a logical value, whether to connect the
#' nodes of the same group with wires.
#' @param shape a logical value, whether to use nodes of
#' different shapes to distinguish the group it belongs to
#' @param facet apply `facet_grid` to the plot by specified variable, e.g., "ONTOLOGY", "category" and "intersect".
#' @param strip_width width of strip text, a.k.a facet label.
#' @param colorBy variable that used to color enriched terms,
#' e.g. 'pvalue', 'p.adjust' or 'qvalue'
#' @importFrom ggplot2 facet_grid
#' @importFrom ggplot2 scale_size_continuous
#' @importFrom rlang check_installed  
dotplot.compareClusterResult <- function(object, x= "Cluster", colorBy="p.adjust",
                                         showCategory=5, by="geneRatio", size="geneRatio",
                                         split=NULL, includeAll=TRUE,
                                         font.size=12, title="", label_format = 30,
                                         group = FALSE, shape = FALSE, facet=NULL, strip_width=15) {                                   
    color <- NULL
    if (is.null(size)) size <- by ## by may deprecated in future release

    if (!is.null(facet) && facet == "intersect") {
        object <- append_intersect(object)
    }

    df <- fortify(object, showCategory=showCategory, by=size,
                  includeAll=includeAll, split=split)

    # if (by != "geneRatio")
    #    df$GeneRatio <- parse_ratio(df$GeneRatio)
    label_func <- default_labeller(label_format)
    if(is.function(label_format)) {
        label_func <- label_format
    }

    if (size %in% c("rowPercentage", "count", "geneRatio")) {
        by2 <- switch(size, rowPercentage = "Percentage",
                            count         = "Count",
                            geneRatio     = "GeneRatio")
    } else {
        by2 <- size
    }

    p <- ggplot(df, aes_string(x = x, y = "Description", size = by2)) +
        scale_y_discrete(labels = label_func)

    if (group) {
        p <- p + geom_line(aes_string(color = "Cluster", group = "Cluster"), size=.3) +
          ggnewscale::new_scale_colour()
    }


    if (shape) {
        check_installed('ggstar', 'for `dotplot()` with `shape = TRUE`.')
        ggstar <- "ggstar"
        require(ggstar, character.only=TRUE)
        # p <- p + ggsymbol::geom_symbol(aes_string(symbol = "Cluster", fill = colorBy)) +
        p <- p + ggstar::geom_star(aes_string(starshape="Cluster", fill=colorBy)) +
            set_enrichplot_color(type = "fill")
    }  else {
        p <- p +  geom_point(aes_string(fill = colorBy)) + 
            aes(shape = I(enrichplot_point_shape))
    }

    p <- p + set_enrichplot_color(type = "fill") +
        ylab(NULL) + ggtitle(title) + DOSE::theme_dose(font.size) +
        scale_size_continuous(range=c(3, 8))


    if (!is.null(facet)) {
        p <- p + facet_grid(.data[[facet]] ~ ., 
                scales = "free", space = 'free',
                switch = 'y',
                labeller = ggplot2::label_wrap_gen(strip_width)) +
            theme(strip.text = element_text(size = 14))
    }

    class(p) <- c("enrichplotDot", class(p))
    
    return(p)
}


append_intersect <- function(x) {
    if (!inherits(x, 'compareClusterResult')) stop("x should be a compareClusterResult object")

    d <- as.data.frame(x)
    # sets <- split(d$Description, d$Cluster)
    # yulab.utils::check_pkg('aplotExtra', 'for upsetplot of compareCluster Result')
    # tidy_main_subsets <- yulab.utils::get_fun_from_pkg('aplotExtra', 'tidy_main_subsets')

    # df <- tidy_main_subsets(sets, order.intersect.by = 'name', nintersects = 10, order.set.by = 'name')
    # nn <- names(sets)
    # df2 <- tidyr::unnest(df[, c('id', 'item')], cols="item")
    # so <- vapply(df2$id, function(x) {
    #         paste(nn[as.numeric(strsplit(x, '/')[[1]])], collapse = " & ")
    #     }, character(1)
    # )


    # set_info <- data.frame(
    #     intersect = so, 
    #     Description = df2$item
    # )
    

    so <- vapply(split(d$Cluster, d$Description), 
        FUN = paste, 
        FUN.VALUE = character(1),
        collapse = " & ")

    set_info <- data.frame(
        intersect = so, 
        Description = names(so)
    )

    d2 <- merge(d, set_info, by="Description")
    n <- levels(d2$Cluster)
    cc <- yulab.utils::combinations(length(n))
    lv <- vapply(cc, function(i) paste(n[i], collapse = " & "), character(1))
    d2$intersect <- factor(d2$intersect, levels=lv)
    
    x@compareClusterResult <- d2

    return(x)
}

#' compare two clusters in the compareClusterResult object
#'
#'
#' @title dotplot2
#' @param object a compareClusterResult object
#' @param x selected variable to visualize in x-axis
#' @param vars selected Clusters to be compared, only length of two is supported
#' @param label to label the Clusters in the plot, should be a named vector
#' @param ... additional parameters passed to dotplot
#' @return a ggplot object
#' @importFrom ggplot2 geom_segment
#' @importFrom ggplot2 geom_vline
#' @importFrom ggplot2 geom_blank
#' @importFrom grid arrow
#' @importFrom grid unit
#' @importFrom ggplot2 geom_text
#' @export
#' @author Guangchuang Yu
dotplot2 <- function(object, x="FoldEnrichment", vars = NULL, label = "auto", ...) {
    if (!is(object, 'compareClusterResult')) {
        stop('only compareClusterResult object is supported')
    }

    if (length(vars) != 2) {
        stop("vars should be of length 2.")
    }
    object <- dplyr::filter(object, .data$Cluster %in% vars)
    d <- object@compareClusterResult
    d[[x]] <- d[[x]] * ifelse(d$Cluster == vars[1], -1, 1)
    object@compareClusterResult <-  d
    p <- dotplot(object, x = x, ...)
    p <- p + geom_segment(aes(xend=0, yend=.data$Description), 
            size=1, color='grey50') + 
        geom_vline(xintercept=0, lty='dashed') + 
        scale_x_continuous(labels = abs)
    
    p$layers <- rev(p$layers)
    
    if (is.null(label)) return (p)

    d <- dplyr::group_by(object, .data$Cluster) |> 
        dplyr::summarise(mid = max(abs(.data[[x]])) * sign(max(.data[[x]]))/2)

    if (label != "auto") {
        d$Cluster <- label[d$Cluster]
    }

    p + geom_text(aes(x=.data$mid, y=.4, label=.data$Cluster), 
                data=d, inherit.aes = FALSE, size=5) + 
        geom_segment(aes(x=0, xend=.data$mid, y=.2, yend=.2), 
                data=d, inherit.aes=FALSE, 
                arrow = arrow(length=unit(0.30,"cm"), type = "closed")) +
        geom_blank(data = data.frame(y=0), aes(y=.data$y), inherit.aes=FALSE)

}