R/plot.R

In primerTree: Visually Assessing the Specificity and Informativeness of Primer Pairs

#PrimerTree
#Copyright (C) 2013 Jim Hester

#' plots a tree along with a series of taxonomic ranks
#' @param ranks The ranks to include, defaults to all common ranks, if null
#' print all ranks.
#' @inheritParams plot_tree
#' @seealso \code{\link{plot_tree}} to plot only a single rank or the just the
#' tree layout.
#' @export
#' @examples
#' library(gridExtra)
#' library(directlabels)
#' #plot all the common ranks
#' plot_tree_ranks(mammals_16S$tree, mammals_16S$taxonomy)
#' #plot specific ranks, with a larger dot size
#' plot_tree_ranks(mammals_16S$tree, mammals_16S$taxonomy,
#'   ranks=c('kingdom', 'class', 'family'), size=3)

plot_tree_ranks <- function(
    tree,
    taxonomy,
    main = NULL,
    type = "unrooted",
    ranks = common_ranks,
    size = 2,
    guide_size = NULL,
    legend_cutoff = 25,
    ...
) {
    if (is.null(ranks)) {
        ranks <- setdiff(names(taxonomy), c("accession", "taxId"))
    }

    plots <- list()
    plots$structure <- plot_tree(
        tree,
        main = main,
        guide_size = guide_size,
        type = type,
        ...
    )

    for (rank in intersect(ranks, names(taxonomy))) {
        if (length(na.omit(taxonomy[rank])) > 0) {
            plots[[rank]] <- plot_tree(
                tree,
                guide_size = guide_size,
                type = type,
                rank = rank,
                taxonomy = taxonomy,
                size = size,
                legend_cutoff = legend_cutoff,
                ...
            )
        }
    }
    p <- do.call(gridExtra::arrangeGrob, plots)
    class(p) <- c("primerTree_plot_multi", class(p))

    p
}

common_ranks <- c(
    "kingdom",
    "phylum",
    "class",
    "order",
    "family",
    "genus",
    "species"
)

#' plots a tree, optionally with colored and labeled points by taxonomic rank
#'
#' @param tree to be plotted, use layout_tree to layout tree.
#' @param taxonomy A data.frame with an accession field corresponding to the
#'        tree tip labels.
#' @param main An optional title for the plot
#' @param type The type of tree to plot, default unrooted.
#' @param rank The rank to include, if null only the tree is plotted
#' @param size The size of the colored points
#' @param guide_size The size of the length guide.  If NULL auto detects a
#'        reasonable size.
#' @param legend_cutoff The number of different taxa names after which the
#'        names are no longer printed.
#' @param ... additional arguments passed to \code{\link{layout_tree_ape}}
#' @return plot to be printed.
#' @export

plot_tree <- function(
    tree,
    type = "unrooted",
    main = NULL,
    guide_size = NULL,
    rank = NULL,
    taxonomy = NULL,
    size = 2,
    legend_cutoff = 25,
    ...
) {
    x <- layout_tree_ape(tree, type = type, ...)

    range_x <- range(x$edge$x, x$tip$x)
    range_y <- scales::expand_range(range(x$edqe$y, x$tip$y), mul = .1)

    if (is.null(guide_size)) {
        guide_size <- 10**(round_any(log10(range_x[2] - range_x[1]), 1) - 1)
    }
    p <- ggplot2::ggplot() +
        ggplot2::geom_segment(
            data = x$edge,
            ggplot2::aes_string(x = "x", y = "y", xend = "xend", yend = "yend")
        ) +
        theme_noaxis() +
        ggplot2::annotate(
            "segment",
            x = range_x[1],
            xend = range_x[1] + guide_size,
            y = range_y[1],
            yend = range_y[1],
            arrow = ggplot2::arrow(
                ends = "both",
                angle = 90,
                length = unit(.2, "cm")
            )
        ) +
        ggplot2::annotate(
            "text",
            x = range_x[1] + (guide_size / 2),
            y = range_y[1],
            label = guide_size,
            vjust = -.5
        )

    if (!is.null(rank)) {
        if (is.null(taxonomy)) {
            stop("Must provide a taxonomy if plotting a rank")
        }

        if (is.null(main)) {
            main <- rank
        }

        x$tip <- merge(
            x$tip,
            taxonomy,
            by.x = "label",
            by.y = "accession",
            all.x = TRUE
        )

        rows <- na.omit(x$tip[, c("x", "y", rank)])
        p <- p +
            ggplot2::geom_point(
                data = rows,
                ggplot2::aes_string(x = "x", y = "y", color = rank),
                size = size,
                na.rm = TRUE
            ) +
            ggplot2::theme(legend.position = "none")

        if (length(unique(rows[[rank]])) < legend_cutoff) {
            smart_grid2 <- list("get.means", "calc.boxes", "empty.grid")
            p <- p +
                directlabels::geom_dl(
                    data = rows,
                    method = smart_grid2,
                    ggplot2::aes_string(
                        x = "x",
                        y = "y",
                        color = rank,
                        label = rank
                    )
                )
        }
    }
    p <- p + ggplot2::ggtitle(main)
    class(p) <- c(class(p), "primerTree_plot")
    p
}

#' layout a tree using ape, return an object to be plotted by
#' \code{\link{plot_tree}}
#' @param tree The \code{\link[ape]{phylo}} tree to be plotted
#' @param ... additional arguments to \code{\link[ape]{plot.phylo}}
#' @return \item{edge}{list of x, y and xend, yend coordinates
#' as well as ids for the edges}
#' \item{tips}{list of x, y, label and id for the tips}
#' \item{nodes}{list of x, y and id for the nodes}
layout_tree_ape <- function(tree, ...) {
    #hack to write no output
    cur_dev <- dev.cur() #store previous dev
    temp_file <- tempfile()
    pdf(file = temp_file)
    ape::plot.phylo(tree, plot = FALSE, ...)
    dev.off()
    unlink(temp_file)
    dev.set(cur_dev) #restore previous dev

    last <- .PlotPhyloEnv$last_plot.phylo
    new <- list()
    new$edge$x <- last$xx[last$edge[, 1]]
    new$edge$xend <- last$xx[last$edge[, 2]]

    new$edge$y <- last$yy[last$edge[, 1]]
    new$edge$yend <- last$yy[last$edge[, 2]]
    new$edge$id <- tree$edge

    new$edge <- data.frame(new$edge, stringsAsFactors = FALSE)

    new$tip$x <- last$xx[1:last$Ntip]
    new$tip$y <- last$yy[1:last$Ntip]
    new$tip$label <- tree$tip.label
    new$tip$id <- 1:last$Ntip

    new$tip <- data.frame(new$tip, stringsAsFactors = FALSE)

    new$node$x <- last$xx[(last$Ntip + 1):length(last$xx)]
    new$node$y <- last$yy[(last$Ntip + 1):length(last$yy)]
    new$node$label <- tree$node.label
    new$node <- data.frame(new$node, stringsAsFactors = FALSE)

    new
}
theme_noaxis <- function() {
    ggplot2::theme(
        panel.border = ggplot2::element_blank(),
        panel.grid = ggplot2::element_blank(),
        axis.line = ggplot2::element_blank(),
        axis.text = ggplot2::element_blank(),
        axis.title = ggplot2::element_blank(),
        axis.ticks = ggplot2::element_blank(),
        plot.margin = ggplot2::unit(c(0, 0, -1, -1), "lines")
    )
}

#' @export
print.primerTree_plot_multi <- function(x, ...) {
    gridExtra::grid.arrange(x, ...)
}