R/genetracks_grob.R

In myles-lewis/locuszoomr: Gene Locus Plot with Gene Annotations

#' Create gene tracks grob
#' 
#' Plot gene annotation tracks from `ensembldb` data using the grid package to
#' create a grob.
#' 
#' @details This function is called by [gg_genetracks()]. It can be used to
#'   generate a grob of the gene annotation tracks on their own.
#' @param locus Object of class 'locus' generated by [locus()].
#' @param filter_gene_name Vector of gene names to display.
#' @param filter_gene_biotype Vector of gene biotypes to be filtered. Use
#' [ensembldb::listGenebiotypes()] to display possible biotypes. For example, 
#' `ensembldb::listGenebiotypes(EnsDb.Hsapiens.v75)`
#' @param cex.text Font size for gene text.
#' @param showExons Logical whether to show exons or simply show whole gene as a
#'   rectangle. If `showExons = FALSE` colours are specified by `exon_border`
#'   for rectangle border and `gene_col` for the fill colour.
#' @param maxrows Specifies maximum number of rows to display in gene 
#' annotation panel.
#' @param border Logical whether a bounding box is plotted.
#' @param gene_col Colour for gene lines.
#' @param exon_col Fill colour for exons.
#' @param exon_border Border line colour outlining exons (or genes if
#'   `showExons` is `FALSE`). Set to `NA` for no border.
#' @param text_pos Character value of either 'top' or 'left' specifying 
#' placement of gene name labels.
#' @param highlight Vector of genes to highlight.
#' @param highlight_col Single colour or vector of colours for highlighted
#'   genes.
#' @param blanks Controls handling of genes with blank names: `"fill"` replaces
#'   blank gene symbols with ensembl gene ids. `"hide"` hides genes which are
#'   missing gene symbols.
#' @return A grob object.
#' @examples
#' if(require(EnsDb.Hsapiens.v75)) {
#' data(SLE_gwas_sub)
#' loc <- locus(SLE_gwas_sub, gene = 'IRF5', flank = c(7e4, 2e5), LD = "r2",
#'              ens_db = "EnsDb.Hsapiens.v75")
#' g <- genetracks_grob(loc)
#' grid::grid.newpage()
#' grid::grid.draw(g)
#' }
#' @importFrom grid viewport rectGrob textGrob xaxisGrob gList gTree gpar
#'   polylineGrob
#' @export

genetracks_grob <- function(locus,
                            filter_gene_name = NULL,
                            filter_gene_biotype = NULL,
                            border = FALSE,
                            cex.text = 0.7,
                            gene_col = ifelse(showExons, 'blue4', 'skyblue'),
                            exon_col = 'blue4',
                            exon_border = 'blue4',
                            showExons = TRUE,
                            maxrows = NULL,
                            text_pos = 'top',
                            highlight = NULL,
                            highlight_col = "red",
                            blanks = c("fill", "hide")) {
  if (!inherits(locus, "locus")) stop("Object of class 'locus' required")
  blanks <- match.arg(blanks)
  TX <- locus$TX
  EX <- locus$EX
  xrange <- locus$xrange
  if (!is.null(filter_gene_name)) {
    TX <- TX[TX$gene_name %in% filter_gene_name, ]
  }
  if (!is.null(filter_gene_biotype)) {
    TX <- TX[TX$gene_biotype %in% filter_gene_biotype, ]
  }
  if (nrow(TX) == 0) {
    message('No genes to plot')
    return(invisible(NULL))
  }
  
  TX <- gene_colours(TX, gene_col, exon_col, exon_border, showExons,
                     highlight, highlight_col)
  
  TX <- mapRow(TX, xlim = xrange, cex.text = cex.text, text_pos = text_pos,
               blanks = blanks)
  maxrows <- if (is.null(maxrows)) max(TX$row) else min(c(max(TX$row), maxrows))
  if (max(TX$row) > maxrows) message(max(TX$row), " tracks needed to show all genes")
  TX <- TX[TX$row <= maxrows, ]
    
  ylim <- c(-maxrows - 0.3, -0.3)
  xrange <- xrange / 1e6
  TX$start <- TX$start / 1e6
  TX$end <- TX$end / 1e6
  TX[, c("mean", "tmin", "tmax", "min", "max")] <- TX[, c("mean", "tmin", "tmax", "min", "max")] / 1e6
  exheight <- switch(text_pos, "top" = 0.15, "left" = 0.3)
  
  gt <- gTree(
    childrenvp = genetrack.vp(xrange, ylim),
    children = gList(
      if (border) rectGrob(gp = gpar(lwd = 1.5), vp = "genetrack"),
      exonGrob(TX, EX, showExons, exheight),
      genetextGrob(text_pos, TX, xrange, cex.text)),
    gp = gpar()
  )
  
  gt
}


genetrack.vp <- function(xrange, ylim) {
  viewport(name = "genetrack",
           x = unit(0, "lines"),
           y = unit(0, "lines"),
           width = unit(1, "npc"),
           height = unit(1, "npc") - unit(0, "lines"),
           just = c("left", "bottom"),
           xscale = xrange + c(-0.04, 0.04) * diff(xrange),
           yscale = ylim)
}


exonGrob <- function(TX, EX, showExons, exheight) {
  if (showExons) {
    LX <- unlist(t(TX[, c('start', 'end')]))
    LY <- cbind(-TX[, 'row'], -TX[, 'row'])
    LY <- unlist(t(LY))
    line_id <- rep(seq_len(nrow(TX)), each = 2)
    
    EXset <- lapply(seq_len(nrow(TX)), function(i) {
      e <- EX[EX$gene_id == TX$gene_id[i], ]
      exstart <- start(e) / 1e6
      exwidth <- end(e) / 1e6 - exstart
      data.frame(x = exstart,
                 y = -TX[i, 'row'] - exheight,
                 width = exwidth,
                 height = 2 * exheight,
                 exon_col = TX$exon_col[i],
                 exon_border = TX$exon_border[i])
    })
    EXset <- do.call("rbind", EXset)
    
    gList(
      polylineGrob(unit(LX, "native"),
                   unit(LY, "native"),
                   id = line_id,
                   gp = gpar(col = TX$gene_col, lwd = 1.5, lineend = "butt"),
                   vp = "genetrack"),
      rectGrob(x = unit(EXset[, "x"], "native"),
               y = unit(EXset[, "y"], "native"),
               width = unit(EXset[, "width"], "native"),
               height = unit(EXset[, "height"], "native"),
               just = c("left", "bottom"),
               gp = gpar(fill = EXset$exon_col, col = EXset$exon_border,
                         lwd = 0.5, lineend = "square", linejoin = "mitre"),
               vp = "genetrack")
    )
  } else {
    # without exons
    rectGrob(x = unit(TX$start, "native"),
             y = unit(-TX[, 'row'] - exheight, "native"),
             width = unit(TX$end - TX$start, "native"),
             height = unit(exheight*2, "native"),
             just = c("left", "bottom"),
             gp = gpar(fill = TX$gene_col, col = TX$exon_border,
                       lineend = "square", linejoin = "mitre"),
             vp = "genetrack")
  }
}


genetextGrob <- function(text_pos, TX, xrange, cex.text) {
  if (text_pos == "top") {
    tfilter <- which(TX$tmin > (xrange[1] - diff(xrange) * 0.04) & 
                       (TX$tmax < xrange[2] + diff(xrange) * 0.04))
    tg <- lapply(tfilter, function(i) {
      textGrob(label = if (TX$strand[i] == "+") {
        bquote(.(TX$gene_name[i]) * symbol("\256"))
      } else {     
        bquote(symbol("\254") * .(TX$gene_name[i]))
      },
      x = unit(TX$mean[i], "native"),
      y = unit(-TX$row[i] + 0.45, "native"),
      gp = gpar(cex = cex.text), vp = "genetrack")
    })
  } else if (text_pos == "left") {
    tfilter <- which(TX$tmin > xrange[1])
    tg <- lapply(tfilter, function(i) {
      textGrob(label = if (TX$strand[i] == "+") {
        bquote(.(TX$gene_name[i]) * symbol("\256"))
      } else {     
        bquote(symbol("\254") * .(TX$gene_name[i]))
      },
      x = unit(pmax(TX$start[i],
                    xrange[1] - diff(xrange) * 0.04) - diff(xrange) * 0.01,
               "native"),
      y = unit(-TX$row[i], "native"),
      just = "right",
      gp = gpar(cex = cex.text), vp = "genetrack")
    })
  }
  do.call(gList, tg)
}