R/alex.waterfall.plot.R

In GRIN2: Genomic Random Interval (GRIN)

#' Generate Waterfall Plot of Lesion and Expression Data
#'
#' @description
#' Generates a waterfall plot displaying genomic lesions and gene expression
#' levels across subjects for a selected gene. Subjects are grouped according
#' to lesion status and ordered by expression level within each lesion group.
#'
#' @usage
#' alex.waterfall.plot(
#'   waterfall.prep,
#'   lsn.data,
#'   lsn.clrs = NULL,
#'   delta = 0.5
#' )
#'
#' @param waterfall.prep Output from \code{\link{alex.waterfall.prep}}. A list
#' containing \code{"gene.lsn.exp"} with subject IDs, lesion groups, and
#' expression values for the selected gene; \code{"lsns"} with genomic lesions
#' overlapping the gene; \code{"stats"} with the corresponding Kruskal-Wallis
#' lesion-expression association results; and \code{"gene.ID"} with the gene
#' symbol or Ensembl gene ID used to label the plot.
#'
#' @param lsn.data A data frame containing genomic lesion data in GRIN-compatible
#' format. It must contain the columns \code{"ID"}, \code{"chrom"},
#' \code{"loc.start"}, \code{"loc.end"}, and \code{"lsn.type"}.
#'
#' @param lsn.clrs Optional named vector specifying colors for lesion groups.
#' Names must correspond to lesion types represented in the data. If
#' \code{NULL}, colors for individual lesion types are assigned using
#' \code{\link{default.grin.colors}}, with additional colors assigned to
#' \code{"none"} and \code{"multiple"} groups.
#'
#' @param delta Numeric value controlling the genomic spacing around the gene
#' locus displayed in the DNA lesion panel. The default is \code{0.5}.
#'
#' @details
#' The left portion of the waterfall plot displays genomic lesions overlapping
#' the selected gene, with lesion types distinguished by color. The genomic
#' coordinates of the gene are indicated by vertical reference lines.
#'
#' The right portion displays gene expression for the same subjects. Subjects
#' are first grouped alphabetically according to lesion group and then ordered
#' by expression level within each group. For each subject, expression is shown
#' relative to the median expression of the corresponding lesion group.
#'
#' Colors may be supplied through \code{lsn.clrs}. When colors are not supplied,
#' lesion-specific colors are assigned automatically using
#' \code{\link{default.grin.colors}}.
#'
#' @return
#' Generates a waterfall plot showing genomic lesion status and gene expression
#' for the selected gene.
#'
#' @export
#'
#' @importFrom stats median aggregate
#' @importFrom graphics plot rect legend text segments
#'
#' @references
#' Cao, X., Elsayed, A. H., & Pounds, S. B. (2023). Statistical Methods Inspired
#' by Challenges in Pediatric Cancer Multi-omics.
#'
#' @author
#' Abdelrahman Elsayed \email{abdelrahman.elsayed@stjude.org},
#' Stanley Pounds \email{stanley.pounds@stjude.org}
#'
#' @seealso
#' \code{\link{alex.prep.lsn.expr}},
#' \code{\link{KW.hit.express}},
#' \code{\link{alex.waterfall.prep}}
#'
#' @examples
#' data(expr_data)
#' data(lesion_data)
#' data(hg38_gene_annotation)
#'
#' # Prepare expression and lesion data
#' alex.data <- alex.prep.lsn.expr(expr_data,
#'                                 lesion_data,
#'                                 hg38_gene_annotation,
#'                                 min.expr = 1,
#'                                 min.pts.lsn = 5)
#'
#' # Run Kruskal-Wallis test
#' alex.kw.results <- KW.hit.express(alex.data,
#'                                   hg38_gene_annotation,
#'                                   min.grp.size = 5)
#'
#' # Prepare data for the WT1 gene
#' WT1.waterfall.prep <- alex.waterfall.prep(alex.data,
#'                                           alex.kw.results,
#'                                           "WT1",
#'                                           lesion_data)
#'
#' # Generate waterfall plot for WT1
#' alex.waterfall.plot(WT1.waterfall.prep,
#'                     lesion_data)
alex.waterfall.plot=function(waterfall.prep,   # Output of the alex.waterfall.prep function
                             lsn.data,         # Lesion data in a GRIN compatible format
                             lsn.clrs=NULL,    # Colors assigned for each lesion group. If NULL, the default.grin.colors function will be used to assign lesion colors automatically
                             delta=0.5)        # spacing argument for the waterfall plot

{
  # Validate input data

  if (!is.list(waterfall.prep))
    stop("waterfall.prep must be the output from alex.waterfall.prep().")

  required.prep.objects=c("gene.lsn.exp","lsns","stats","gene.ID")
  if (!all(required.prep.objects %in% names(waterfall.prep)))
    stop("waterfall.prep must contain: gene.lsn.exp, lsns, stats, and gene.ID.")

  if (!is.data.frame(lsn.data))
    stop("lsn.data must be a data frame.")

  required.lsn.cols=c("ID","chrom","loc.start","loc.end","lsn.type")
  if (!all(required.lsn.cols %in% colnames(lsn.data)))
    stop("lsn.data must contain the columns: ",paste(required.lsn.cols,collapse=", "), ".")

  if (!is.numeric(delta) || length(delta)!=1 || is.na(delta) || delta<0)
    stop("delta must be a single non-negative numeric value.")

  unique.grps=sort(unique(lsn.data$lsn.type))

  # assign colors for multiple and none lesion groups (colors to be assigned automatically for other lesion groups)

  if (is.null(lsn.clrs))
  {
    lsn.grps.clr=default.grin.colors(unique.grps)
    common.grps.clr=c(none="gray", multiple="violet")
    lsn.clrs=c(lsn.grps.clr, common.grps.clr)
  }

  gene.ID=waterfall.prep$gene.ID
  gene.lsn.exp=waterfall.prep$gene.lsn.exp

  lsn.clm=paste0(gene.ID,".lsn")
  rna.clm=paste0(gene.ID,".RNA")

  if (!all(c(lsn.clm,rna.clm) %in% colnames(gene.lsn.exp)))
    stop("waterfall.prep$gene.lsn.exp does not contain the expected lesion and expression columns.")

  plot.grps=unique(gene.lsn.exp[,lsn.clm])

  if (!all(plot.grps %in% names(lsn.clrs)))
    stop("lsn.clrs must provide colors for all lesion groups represented in waterfall.prep.")

  if (all(is.na(gene.lsn.exp[,rna.clm])))
    stop("No expression values are available for the selected gene.")

  if (diff(range(gene.lsn.exp[,rna.clm],na.rm=TRUE))==0)
    stop("Expression values for the selected gene have no variation and cannot be displayed in the waterfall plot.")

  gene.lsn.exp.ord=order(gene.lsn.exp[,lsn.clm],
                         gene.lsn.exp[,rna.clm])

  gene.lsn.exp=gene.lsn.exp[gene.lsn.exp.ord,]

  pt.num=1:nrow(gene.lsn.exp)
  names(pt.num)=gene.lsn.exp$ID

  ####################################

  # Set up plotting region

  plot(c(-1.1,+1.5),
       c(0.1,-1.1)*nrow(gene.lsn.exp),
       type="n",axes=F,
       xlab="",ylab="")

  ###################################

  # DNA lesion plot
  # gene locus

  loc.rng=c(waterfall.prep$stats$loc.start,
            waterfall.prep$stats$loc.end)
  loc.lng=diff(loc.rng)

  pos.rng=loc.rng+c(-1,1)*delta*loc.lng
  waterfall.prep$lsns$x.start=(waterfall.prep$lsns$loc.start-pos.rng[1])/(diff(pos.rng))-1.1
  waterfall.prep$lsns$x.end=(waterfall.prep$lsns$loc.end-pos.rng[1])/(diff(pos.rng))-1.1

  x.locus=(loc.rng-pos.rng[1])/diff(pos.rng)-1.1

  # background for lesion plot

  graphics::rect(-1.1,-nrow(gene.lsn.exp),
                 -0.1,0,
                 col=lsn.clrs["none"],
                 border=lsn.clrs["none"])

  waterfall.prep$lsns$size=waterfall.prep$lsns$loc.end-
    waterfall.prep$lsns$loc.start+1
  ord=rev(order(waterfall.prep$lsns$size))
  waterfall.prep$lsns=waterfall.prep$lsns[ord,]
  graphics::rect(pmax(waterfall.prep$lsns$x.start,-1.1),
                 -pt.num[waterfall.prep$lsns$ID],
                 pmin(waterfall.prep$lsns$x.end,-0.1),
                 -pt.num[waterfall.prep$lsns$ID]+1,
                 col=lsn.clrs[waterfall.prep$lsns$lsn.type],
                 border=lsn.clrs[waterfall.prep$lsns$lsn.type])

  graphics::segments(x.locus,-nrow(gene.lsn.exp),
                     x.locus,0,col="white",lty=3)

  graphics::text(x.locus,-1.05*nrow(gene.lsn.exp),
                 loc.rng,cex=0.75)
  graphics::text(-0.55,+0.05*nrow(gene.lsn.exp),
                 paste0(gene.ID," DNA Lesions"))

  ##############################################

  # RNA expression plot

  rna.lbls=pretty(gene.lsn.exp[,rna.clm])
  ok.lbls=(rna.lbls>min(gene.lsn.exp[,rna.clm],na.rm=T))&
    (rna.lbls<max(gene.lsn.exp[,rna.clm],na.rm=T))
  rna.lbls=rna.lbls[ok.lbls]

  gene.lsn.exp$x.rna=(gene.lsn.exp[,rna.clm]-min(gene.lsn.exp[,rna.clm],na.rm=T))/diff(range(gene.lsn.exp[,rna.clm],na.rm=T))
  gene.lsn.exp$x.rna=gene.lsn.exp$x.rna+0.1

  x.rna.lbls=(rna.lbls-min(gene.lsn.exp[,rna.clm],na.rm=T))/diff(range(gene.lsn.exp[,rna.clm],na.rm=T))+0.1

  n.mdn=function(x)
  {
    mdn=stats::median(x,na.rm=T)
    n=sum(!is.na(x))
    res=unlist(c(n=n,mdn=mdn))
    return(res)
  }

  lsn.mdn=stats::aggregate(x=gene.lsn.exp$x.rna,
                           by=list(lsn=gene.lsn.exp[,lsn.clm]),
                           FUN=n.mdn)
  rownames(lsn.mdn$x)=lsn.mdn$lsn

  x.mdn=lsn.mdn$x[gene.lsn.exp[,lsn.clm],"mdn"]

  graphics::text(x.rna.lbls,-1.05*nrow(gene.lsn.exp),
                 rna.lbls,cex=0.75)
  graphics::segments(x.rna.lbls,0,
                     x.rna.lbls,-nrow(gene.lsn.exp),
                     lty=3)

  graphics::rect(pmin(gene.lsn.exp$x.rna,x.mdn),-(1:nrow(gene.lsn.exp)),
                 pmax(gene.lsn.exp$x.rna,x.mdn),-(1:nrow(gene.lsn.exp))+1,
                 col=lsn.clrs[gene.lsn.exp[,lsn.clm]],
                 border=NA)

  graphics::segments(x.mdn,-(1:nrow(gene.lsn.exp)),
                     x.mdn,-(1:nrow(gene.lsn.exp))+1,
                     col=lsn.clrs[gene.lsn.exp[,lsn.clm]])

  graphics::text(0.55,0.05*nrow(gene.lsn.exp),
                 paste0(gene.ID," RNA Expression"))

  #############################

  # Add legend

  lsn.inc=names(lsn.clrs)%in%gene.lsn.exp[,lsn.clm]
  graphics::legend(1.05,-0.25*nrow(gene.lsn.exp),
                   fill=unlist(lsn.clrs[lsn.inc]),
                   legend=names(lsn.clrs[lsn.inc]),
                   cex=0.75,border=NA,bty="n")

}