R/Plot_functions.R

In CL-CHEN-Lab/User_interface_for_Kronos_scRT: R-package and Shiny App interface for the analysis of single cell replication timing data

#'  scRT plot of a region of interest
#'
#' @return plot
#'
#' @importFrom  dplyr tibble filter rename arrange mutate summarise ungroup
#' @importFrom  tidyr %>% gather
#' @importFrom  ggplot2 aes annotate element_text facet_grid geom_path geom_rect ggplot labs scale_color_manual scale_fill_gradient scale_fill_manual scale_x_continuous scale_y_continuous sec_axis theme theme_bw theme_set xlab
#' @importFrom ggrastr geom_tile_rast
#' @importFrom stringr str_count
#' @importFrom RColorBrewer brewer.pal
#'
#' @param pseudoBulkRT, pseudo-Bulk RT that can be rebinned with a reference RT (generated by RebinRT)
#' @param S_scCN, S-phase scCN dataframe created by Replication_state (optionally filtered)
#' @param Coordinates, named list containing chr,start and end position
#' @param rasterized_heatmap, logic: whether or not the heatmap should be rasterized
#' @param heatmap_colors, a vector of colors for the heatmap.(Replicated/min value and Unreplicated/max value)
#' @param sample_colors, a vector, named or not, for the samples
#' @export
#'

scRTplot = function(pseudoBulkRT,
                    S_scCN,
                    Coordinates,
                    Plot = 'scRT',
                    rasterized_heatmap = F,
                    heatmap_colors = NULL,
                    sample_colors = NULL) {
  #load required operators
  `%>%` = tidyr::`%>%`

  #set theme
  ggplot2::theme_set(new = ggplot2::theme_bw())


  #check Plot

  if (!Plot %in% c('scRT', 'S_scCN', 'Norm. S_scCN')) {
    stop('Plot can be either scRT, S_scCN or Norm. S_scCN')
  }


  #calculate extremes
  extremes = S_scCN %>%
    dplyr::ungroup() %>%
    dplyr::summarise(CN_bg = round(stats::quantile(CN_bg, c(0.01, 0.99)), 1),
                     CN = round(stats::quantile(CN, c(0.01, 0.99)), 1))
  #filter data
  pseudoBulkRT = pseudoBulkRT %>%
    dplyr::filter(chr == Coordinates$chr,
                  start >= Coordinates$start,
                  end <= Coordinates$end)
  S_scCN = S_scCN %>%
    dplyr::filter(chr == Coordinates$chr,
                  start >= Coordinates$start,
                  end <= Coordinates$end)

  #depending on selection change aesthetics of the plot
  if (Plot == 'scRT') {
    if (is.null(heatmap_colors)) {
      heatmap_colors = c("Replicated" = '#a7001b',
                         "Unreplicated" = '#005095')
    } else if (any(is.na(names(heatmap_colors)))) {
      names(heatmap_colors) = c("Replicated", "Unreplicated")
    }

    S_scCN = S_scCN %>%
      dplyr::rename(Value = Rep) %>%
      dplyr::mutate(Value = ifelse(Value, "Replicated", "Unreplicated"))

    ggEXTRA = ggplot2::ggplot() +
      ggplot2::scale_fill_manual(values = heatmap_colors) +
      ggplot2::labs(fill = 'State', color = 'Sample')

  } else if (Plot == 'S_scCN') {
    if (is.null(heatmap_colors)) {
      heatmap_colors = c('#c56700',
                         '#6d0042')
    }

    S_scCN = S_scCN %>%
      dplyr::rename(Value = CN)

    ggEXTRA = ggplot2::ggplot() +
      ggplot2::scale_fill_gradient(
        low = heatmap_colors[1],
        high = heatmap_colors[2],
        limits = c(extremes$CN[1], extremes$CN[2]),
        oob = scales::squish
      ) +
      ggplot2::labs(fill = 'CNV ', color = 'Sample')

  } else if (Plot == 'Norm. scCN') {
    if (is.null(heatmap_colors)) {
      heatmap_colors = c('#00c5a4',
                         '#e58225')
    }

    S_scCN = S_scCN %>%
      dplyr::rename(Value = CN_bg)

    ggEXTRA = ggplot2::ggplot() +
      ggplot2::scale_fill_gradient(
        low = heatmap_colors[1],
        high = heatmap_colors[2],
        limits = c(extremes$CN_bg[1], extremes$CN_bg[2]),
        oob = scales::squish
      ) +
      ggplot2::labs(fill = expression(over(S[CNV[i]],
                                           bar(G1 / G2[CNV]))),
                    color = 'Sample')

  }

  Chr = Coordinates$chr
  Maxi = max(S_scCN$newIndex)
  n = stringr::str_count(Maxi)
  basenames_colors = unique(pseudoBulkRT$basename)

  #set colors for basename
  if (is.null(sample_colors)) {
    sample_colors = colorRampPalette(RColorBrewer::brewer.pal(8, "Dark2"))(length(basenames_colors))
    names(sample_colors) = basenames_colors
  } else if (length(sample_colors) < length(basenames_colors) &
             is.null(names(sample_colors))) {
    sample_colors = c(sample_colors,
                      colorRampPalette(RColorBrewer::brewer.pal(8, "Dark2"))(length(basenames_colors) -
                                                                               length(sample_colors)))

    #name the vector
    names(sample_colors) = basenames_colors

  } else if (length(sample_colors) < length(basenames_colors) &
             !any(is.null(names(sample_colors)))) {
    #if the given colors are less than the number of types and the vector is not named
    default_colors = colorRampPalette(RColorBrewer::brewer.pal(8, "Dark2"))(length(basenames_colors) -
                                                                              length(sample_colors))

    default_colors = default_colors[!names(default_colors) %in% names(sample_colors)]

    #merge the two
    sample_colors = c(sample_colors, default_colors)
  }

  plot =
    ggEXTRA +
    ggplot2::geom_path(
      data = pseudoBulkRT %>%
        dplyr::mutate(mid = (start + end) / 2,
                      RT = RT * Maxi / 6 + Maxi / 20) %>%
        tidyr::gather(Plot, pos, start, end) %>%
        dplyr::arrange(mid, pos),
      ggplot2::aes(pos, RT, color = basename)
    ) +
    ggplot2::annotate(
      "rect",
      xmin = -Inf,
      xmax = Inf,
      ymin = Maxi / 100,
      ymax = Maxi / 20 - Maxi / 100,
      fill = 'white'
    )
  if (rasterized_heatmap) {
    plot = plot + ggrastr::geom_tile_rast(data = S_scCN ,
                                          ggplot2::aes(
                                            x = start,
                                            y = -newIndex,
                                            fill = Value
                                          ))


  } else{
    plot = plot +
      ggplot2::geom_rect(
        data = S_scCN ,
        ggplot2::aes(
          xmin = start,
          xmax = end,
          ymin = -newIndex,
          ymax = -newIndex - 1,
          fill = Value
        )
      )
  }

  plot = plot +
    ggplot2::scale_y_continuous(
      breaks = c(Maxi / 6 + Maxi / 20, Maxi / 3 + Maxi / 20, Maxi / 20),
      labels = c('Early - 1', 'Mind - 0.5', 'Late - 0'),
      name = 'RT',
      sec.axis = ggplot2::sec_axis(
        ~ .,
        breaks = c(-seq(1, round(Maxi / 10 ^ (
          n - 1
        )), 1) * 10 ^ (n - 1)) - 0.5,
        labels = as.character(c(seq(
          1, round(Maxi / 10 ^ (n - 1)), 1
        ) * 10 ^ (n - 1))),
        name = 'Single Cell tracks ordered by S-phase progression'
      )
    ) +
    ggplot2::scale_x_continuous(
      labels = function(x)
        paste(x / 10 ^ 6, 'Mb', sep = ' ')
    ) + ggplot2::theme(
      legend.position = 'right',
      legend.direction = "vertical",
      axis.text.x = ggplot2::element_text(
        angle = 45,
        hjust = 1,
        vjust = 1
      ),
      axis.title.y.right = ggplot2::element_text(hjust = 0.6, vjust =
                                                   2),
      axis.title.y.left  = ggplot2::element_text(hjust = 0.92, vjust =
                                                   2)

    ) + ggplot2::xlab(Chr) +
    ggplot2::scale_color_manual(values = sample_colors) +
    ggplot2::facet_grid( ~ group)
  return(plot)

}


#'  scCN plot
#'
#' @return plot
#'
#' @importFrom  dplyr filter select arrange mutate ungroup
#' @importFrom  tidyr %>% drop_na
#' @importFrom  ggplot2 aes element_blank element_text facet_grid geom_rect ggplot scale_fill_manual scale_x_continuous theme theme_bw theme_set
#' @importFrom  foreach %do% foreach
#' @importFrom ggrastr geom_tile_rast
#' @importFrom viridis viridis inferno magma plasma cividis
#'
#' @param S_scCN, S-phase scCN dataframe created by Replication_state (optionally filtered)
#' @param G_scCN, G1/G2-phase scCN dataframe created by Replication_state (optionally filtered)
#' @param CN_limit, Max number of different level to visulize from the lowest one.
#' @param Coordinates, data frame containing chr,start and end position (multiple regions allowed, optional)
#' @param rasterized_heatmap, logic: whether or not the heatmap should be rasterized
#' @param heatmap_color_palette, palette choice (viridis, inferno, magma, plasma or cividis)
#' @param manual_colors, manual colors, if provided they overwrite heatmap_color_palette
#' @export
#'
#'
scCNplot = function(S_scCN,
                    G_scCN,
                    CN_limit = 4,
                    Coordinates = NULL,
                    rasterized_heatmap = F,
                    heatmap_color_palette=c('viridis','inferno','magma','plasma','cividis'),
                    manual_colors=NULL) {
  #load required operators
  `%>%` = tidyr::`%>%`
  `do` = foreach::`%do%`
  if(is.null(manual_colors)){
  if(heatmap_colors[1]=='inferno'){
    heatmap_colors=viridis::inferno(CN_limit + 1)
  }else if(heatmap_colors[1]=='magma'){
    heatmap_colors=viridis::magma(CN_limit + 1)
  }else if(heatmap_colors[1]=='plasma'){
    heatmap_colors=viridis::plasma(CN_limit + 1)
  }else if(heatmap_colors[1]=='cividis'){
    heatmap_colors=viridis::cividis(CN_limit + 1)
  }else{
    heatmap_colors=viridis::viridis(CN_limit + 1)
  }}

  #set theme
  ggplot2::theme_set(new = ggplot2::theme_bw())

  #combine data
  tracks =
    rbind(
      G_scCN %>% dplyr::select(chr, start, end, CN, newIndex, group) %>%
        dplyr::mutate(phase = 'G1G2-phase') %>%
        dplyr::ungroup(),
      S_scCN %>% dplyr::select(chr, start, end, CN, newIndex, group) %>%
        dplyr::mutate(phase = 'S-phase') %>%
        dplyr::ungroup()
    )

  tracks = tracks %>% dplyr::mutate(
    CN = round(CN),
    MaxAllowed = min(CN, na.rm = T) + CN_limit,
    CN = ifelse(CN > MaxAllowed, MaxAllowed, CN)
  ) %>%
    dplyr::arrange(CN) %>%
    dplyr::mutate(CN = factor(CN)) %>%
    tidyr::drop_na() %>%
    dplyr::arrange(chr, start)

  #filter data if provided
  if (!is.null(Coordinates)) {
    tracks = foreach::foreach(i = 1:nrow(Coordinates), .combine = 'rbind') %do%
      {
        tracks %>%
          dplyr::filter(chr == Coordinates$chr[i],
                        start >= Coordinates$start[i],
                        end <= Coordinates$end[i])
      }
  }

  if (rasterized_heatmap) {
    plot =
      ggplot2::ggplot() +
      ggrastr::geom_tile_rast(data = tracks ,
                              ggplot2::aes(
                                x = start,
                                y = -newIndex,
                                fill = CN
                              ))

  } else{
    plot =
      ggplot2::ggplot() +
      ggplot2::geom_rect(
        data = tracks ,
        ggplot2::aes(
          xmin = start,
          xmax = end,
          ymin = -newIndex,
          ymax = -newIndex - 1,
          fill = CN
        )
      )
  }

  plot = plot + ggplot2::facet_grid(phase ~ chr, scales = 'free', space =  'free') +
    ggplot2::scale_x_continuous(
      labels = function(x)
        paste(x / 10 ^ 6, 'Mb', sep = ' ')
    ) + ggplot2::theme(
      legend.position = 'right',
      legend.direction = "vertical",
      axis.title.y = ggplot2::element_blank(),
      axis.text.y = ggplot2::element_blank(),
      axis.ticks.y = ggplot2::element_blank(),
      axis.text.x = ggplot2::element_text(
        angle = 45,
        hjust = 1,
        vjust = 1
      )

    ) +
    ggplot2::scale_fill_manual(values = heatmap_colors)

  return(plot)
}

#'  bulk RT plots of a region of interest
#'
#' @return plot
#'
#' @importFrom  dplyr filter mutate
#' @importFrom  tidyr %>% unnest unite
#' @importFrom  ggplot2 aes annotate geom_line facet_grid theme labs scale_x_continuous theme_bw
#' @importFrom RColorBrewer brewer.pal
#'
#' @param ..., one or multiple dataframes containing bulk and pseudo-bulk RTs and the following columns:chr, start, end, group, basename, RT
#' @param Coordinates, named list containing chr,start and end position
#' @param plotting_groups, optional: if provided samples can be plot together for comparison. It is a named list of arrays where the name is a sample basename if unique, otherwise basename - group, and the array contains the group(s) in which the sample has to be plot.
#' @param sample_colors, a vector, named or not, for each sample. If multiple samples have the same name use "basename - group" otherwise only basename
#' @param highlight_regions, optional: a list of arrays. each array contains start and end position of the region to highlight and optionally the color to use in third position
#'
#' @export

Plot_bulkRT = function(...,
                       Coordinates = list(chr = 'chr1', start = 0, end = Inf),
                       plotting_groups = NULL,
                       sample_colors=NULL,
                       highlight_regions = NULL) {
  #define operator
  `%>%` = tidyr::`%>%`

  #bind data
  data = do.call('rbind', list(...))

  #filter coords
  data = data %>%
    dplyr::filter(chr == Coordinates$chr,
                  start >= Coordinates$start,
                  end <= Coordinates$end) %>%
    dplyr::mutate(mid = (start + end) / 2)

  if(length(unique(data$group))>length(unique(data$basename))){
    data=data%>%
      tidyr::unite(basename,basename,group,sep = ' - ',remove = F)
  }



  if (!is.null(plotting_groups)) {
    for (i in 1:length(plotting_groups)) {
      data$group[data$basename == names(plotting_groups)[i]] = list(plotting_groups[[i]])
    }

    data = data %>%
      tidyr::unnest(cols = c(group))

  }

  basenames_colors=unique(data$basename)
  #set colors for basename
  if (is.null(sample_colors)) {
    sample_colors = colorRampPalette(RColorBrewer::brewer.pal(8, "Dark2"))(length(basenames_colors))
    names(sample_colors) = basenames_colors
  } else if (length(sample_colors) < length(basenames_colors) &
             is.null(names(sample_colors))) {
    sample_colors = c(sample_colors,
                      colorRampPalette(RColorBrewer::brewer.pal(8, "Dark2"))(length(basenames_colors) -
                                                                               length(sample_colors)))

    #name the vector
    names(sample_colors) = basenames_colors

  } else if (length(sample_colors) < length(basenames_colors) &
             !is.null(names(sample_colors))) {
    #if the given colors are less than the number of types and the vector is not named
    default_colors = colorRampPalette(RColorBrewer::brewer.pal(8, "Dark2"))(length(basenames_colors) -
                                                                              length(sample_colors))

    default_colors = default_colors[!names(default_colors) %in% names(sample_colors)]

    #merge the two
    sample_colors = c(sample_colors, default_colors)
  }

  plot = data %>%
    ggplot2::ggplot(ggplot2::aes(mid, RT, color = basename)) +
    ggplot2::geom_line() +
    ggplot2::facet_grid(group ~ .) +
    ggplot2::theme_bw() +
    ggplot2::theme(legend.position = 'top') +
    ggplot2::labs(x = Coordinates$chr, y = 'Replication Timing', color =
                    NULL) +
    ggplot2::scale_x_continuous(
      labels = function(x)
        paste(x / 10 ^ 6, 'Mb', sep = ' ')
    )+
    ggplot2::scale_color_manual(values = sample_colors)

  if (!is.null(highlight_regions)) {
    htoadd = lapply(highlight_regions, function(x) {
      ggplot2::annotate(
        geom = 'rect',
        xmin = as.numeric(x[1]),
        xmax = as.numeric(x[2]),
        ymin = 0,
        ymax = 1,
        fill = ifelse(is.na(x[3]), 'yellow', x[3]),
        alpha = 0.2
      )
    })

    plot = plot + htoadd
  }

  return(plot)
}