R/DIscBIO-generic-PlotmclustMB.R

In DIscBIO: A User-Friendly Pipeline for Biomarker Discovery in Single-Cell Transcriptomics

#' @title Plotting the Model-based clusters in PCA.
#' @description Plot the model-based clustering results
#' @param object \code{DISCBIO} class object.
#' @importFrom ggplot2 ggplot aes geom_point aes_string scale_colour_manual
#'   geom_text geom_segment guides guide_legend xlab ylab theme element_blank
#'   element_line unit element_text element_rect
#' @importFrom igraph get.edgelist degree get.shortest.paths
#' @return A plot of the PCA.
setGeneric("PlotmclustMB", function(object) {
  standardGeneric("PlotmclustMB")
})

#' @export
#' @rdname PlotmclustMB
setMethod(
  "PlotmclustMB",
  signature = "DISCBIO",
  definition = function(object) {
    if (length(object@MBclusters) == 0) {
      stop("run ExprmclustMB before PlotmclustMB")
    }
    total <- object@MBclusters

    Plotmclust <- function(mclustobj, x = 1, y = 2, MSTorder = NULL,
                           show_tree = TRUE, show_full_tree = FALSE,
                           show_cell_names = FALSE, cell_name_size = 3,
                           markerexpr = NULL,
                           showcluster = TRUE) {
      lib_info_with_pseudo <- data.frame(
        State = mclustobj$clusterid,
        sample_name = names(mclustobj$clusterid)
      )
      lib_info_with_pseudo$State <-
        factor(lib_info_with_pseudo$State)
      S_matrix <- mclustobj$pcareduceres
      pca_space_df <- data.frame(S_matrix[, c(x, y)])
      colnames(pca_space_df) <- c("pca_dim_1", "pca_dim_2")
      pca_space_df$sample_name <- row.names(pca_space_df)
      edge_df <- merge(
        pca_space_df, lib_info_with_pseudo,
        by.x = "sample_name", by.y = "sample_name"
      )
      edge_df$Marker <- markerexpr[edge_df$sample_name]
      if (!is.null(markerexpr)) {
        g <- ggplot(
          data = edge_df,
          aes(
            x = pca_space_df$pca_dim_1,
            y = pca_space_df$pca_dim_2,
            size = edge_df$Marker
          )
        )
        if (showcluster) {
          g <- g + geom_point(
            aes_string(color = lib_info_with_pseudo$State),
            na.rm = TRUE
          )
          g <- g + scale_colour_manual(
            values = c(
              "1" = "black",
              "2" = "blue",
              "3" = "green",
              "4" = "red",
              "5" = "yellow",
              "6" = "gray"
            )
          )
        } else {
          g <- g + geom_point(na.rm = TRUE, color = "green")
        }
      } else {
        g <-
          ggplot(
            data = edge_df,
            aes(
              x = pca_space_df$pca_dim_1,
              y = pca_space_df$pca_dim_2
            )
          )
        if (showcluster) {
          g <- g + geom_point(
            aes_string(color = lib_info_with_pseudo$State),
            na.rm = TRUE, size = 3
          )
          g <- g + scale_colour_manual(
            values = c(
              "1" = "black",
              "2" = "blue",
              "3" = "green",
              "4" = "red",
              "5" = "yellow",
              "6" = "gray"
            )
          )
        } else {
          g <- g + geom_point(na.rm = TRUE, size = 3)
        }
      }
      if (show_cell_names) {
        g <- g + geom_text(
          aes(label = names(mclustobj$clusterid)),
          size = cell_name_size
        )
      }

      if (show_tree) {
        clucenter <- mclustobj$clucenter[, c(x, y)]
        clulines <- NULL
        if (show_full_tree) {
          alledges <-
            as.data.frame(get.edgelist(mclustobj$MSTtree),
              stringsAsFactors = FALSE
            )
          alledges[, 1] <- as.numeric(alledges[, 1])
          alledges[, 2] <- as.numeric(alledges[, 2])
          for (i in seq_len(nrow(alledges))) {
            clulines <- rbind(
              clulines,
              c(
                clucenter[alledges[i, 1], ],
                clucenter[alledges[i, 2], ]
              )
            )
          }
        } else {
          if (is.null(MSTorder)) {
            clutable <- table(mclustobj$clusterid)
            alldeg <- degree(mclustobj$MSTtree)
            allcomb <- expand.grid(
              as.numeric(names(alldeg)[alldeg == 1]),
              as.numeric(names(alldeg)[alldeg == 1])
            )
            allcomb <-
              allcomb[allcomb[, 1] < allcomb[, 2], ]
            numres <- t(apply(allcomb, 1, function(i) {
              tmp <- as.vector(
                get.shortest.paths(
                  mclustobj$MSTtree,
                  i[1], i[2]
                )$vpath[[1]]
              )
              c(length(tmp), sum(clutable[tmp]))
            }))
            orderedRows <- order(
              numres[, 1], numres[, 2],
              decreasing = TRUE
            )[1]
            optcomb <- allcomb[orderedRows, ]
            MSTorder <- get.shortest.paths(
              mclustobj$MSTtree, optcomb[1],
              optcomb[2]
            )$vpath[[1]]
          }
          for (i in 1:(length(MSTorder) - 1)) {
            clulines <- rbind(
              clulines,
              c(
                clucenter[MSTorder[i], ],
                clucenter[MSTorder[i + 1], ]
              )
            )
          }
        }
        clulines <- data.frame(
          x = clulines[, 1],
          xend = clulines[, 3],
          y = clulines[, 2],
          yend = clulines[, 4]
        )
        g <- g + geom_segment(
          aes_string(
            x = "x",
            xend = "xend",
            y = "y",
            yend = "yend",
            size = NULL
          ),
          data = clulines,
          size = 1,
          color = "orange"
        )
        clucenter <- data.frame(
          x = clucenter[, 1],
          y = clucenter[, 2],
          id = seq_len(nrow(clucenter))
        )
        g <- g + geom_text(
          aes_string(
            label = "id",
            x = "x",
            y = "y",
            size = NULL
          ),
          data = clucenter,
          size = 10,
          color = "orange"
        )
      }
      g <-
        g +
        guides(
          colour = guide_legend(override.aes = list(size = 5))
        ) +
        xlab(paste0("PCA_dimension_", x)) +
        ylab(paste0("PCA_dimension_", y)) +
        theme(
          panel.border = element_blank(),
          axis.line = element_line()
        ) +
        theme(
          panel.grid.minor.x = element_blank(),
          panel.grid.minor.y = element_blank()
        ) +
        theme(
          panel.grid.major.x = element_blank(),
          panel.grid.major.y = element_blank()
        ) +
        theme(
          legend.position = "top",
          legend.key.size = unit(0.3, "in"),
          legend.text = element_text(size = 20),
          legend.title = element_text(size = 20),
          legend.box = "vertical"
        ) +
        theme(legend.key = element_blank()) +
        theme(panel.background = element_rect(fill = "white")) +
        theme(
          axis.text.x = element_text(size = 17, color = "black"),
          axis.text.y = element_text(size = 17, color = "black"),
          axis.title.x = element_text(size = 20, vjust = -1),
          axis.title.y = element_text(size = 20, vjust = 1),
          plot.margin = unit(c(1, 1, 1, 1), "cm")
        )
      g
    }
    Plotmclust(total)
  }
)