R/plot_distribution_partialTAI.R

In drostlab/myTAI: Evolutionary Transcriptomics

#' @title Comparing partial TAI distributions across stages
#' @description \code{plot_distribution_partialTAI} generates two plots that help to compare the distribution
#' of partial TAI through various developmental stages, highlighting each stage with 
#' distinct colors.
#' @param ExpressionSet a standard PhyloExpressionSet or DivergenceExpressionSet object.
#' @param stages a numeric vector specifying the indices of the stages to compare. Each index 
#' corresponds to a stage in the ExpressionSet. Starts in one.
#' @param xlab label of x-axis.
#' @param ylab label of y-axis.
#' @param main figure title.
#' @param seed defines the colors for the different developmental stages
#' @details
#'  Recommendation - Apply a square root transformation to enhance the visualization of differences 
#' in the distributions: \code{plot_distribution_partialTAI(tf(ExpressionSet, sqrt))}
#' @author Filipa Martins Costa
#' @export

plot_distribution_partialTAI <- function(ExpressionSet,
                                         stages = 1:ncol(ExpressionSet)-2,
                                         xlab = "Partial TAI",
                                         ylab = "Density",
                                         main = "Density Distribution of Partial TAI by Developmental Stage",
                                         seed = 123){
  
  if (any(stages > ncol(ExpressionSet))) {
    stop("Some indices in 'stages' exceed the number of columns in 'ExpressionSet'.")
  }
  
  partial_TAI_matrix <- pMatrix(ExpressionSet[,c(1,2,stages+3)])
  
  partial_TAI_df <- tibble::rownames_to_column(as.data.frame(partial_TAI_matrix), var = "GeneID")
  partial_TAI_long <- tidyr::pivot_longer(
    partial_TAI_df,
    cols = -GeneID,
    names_to = "Stage",
    values_to = "PartialTAI"
  )
  
  qual_col_pals <-  RColorBrewer::brewer.pal.info[ RColorBrewer::brewer.pal.info$category == 'qual',]
  col_vector <- unlist(mapply( RColorBrewer::brewer.pal, qual_col_pals$maxcolors, rownames(qual_col_pals)))
  set.seed(seed)
  colors <- sample(col_vector, ncol(partial_TAI_df))
  
  partial_TAI_long$Stage <- factor(partial_TAI_long$Stage, levels = colnames(partial_TAI_matrix))
  
  P1 <- ggplot2::ggplot(partial_TAI_long, ggplot2::aes(x = PartialTAI, fill = Stage)) +
    ggplot2::geom_density(alpha = 0.7, color = "black") +
    ggplot2::labs(
      x = xlab,
      y = ylab
    ) +
    ggplot2::theme_minimal() +
    ggplot2::theme(
      plot.title = ggplot2::element_text(hjust = 0.5, size = 14),
      axis.title = ggplot2::element_text(size = 12)
    ) +
    ggplot2::scale_fill_manual(values = colors)
  
  P2 <- ggplot2::ggplot(partial_TAI_long, ggplot2::aes(x = PartialTAI, y = Stage, fill = Stage)) +
    ggridges::geom_density_ridges(alpha = 0.7, color = "black", scale = 1) + 
    ggplot2::labs(
      x = xlab,
      y = ylab
    ) +
    ggplot2::theme_minimal() +
    ggplot2::theme(
      plot.title = ggplot2::element_text(hjust = 0.5, size = 14),
      axis.title =ggplot2:: element_text(size = 12)
    ) +
    ggplot2::scale_fill_manual(values = colors)
  
  cowplot::plot_grid(P1, P2, labels = main)
  
}