R/plot.DecisionTreeObject.R

In tomis9/decisionTree: Deciosion tree algorithm

#' plotDecisionTree
#'
#' plot a decision tree
#' @param DecisionTreeObject object of class DecisionTreeObject generated by function decisionTree
#' @param dict a data.frame of optional names for nodes, colnames of data.frame are: vName and toPlot, remember about stringsAsFactors=F
#' @param green which class should be coloured in green? 0 or 1?
#' @param left which class should be mentioned on leaves
#' @keywords plot decision tree, decision tree
#' @export
#' @examples
#' d <- iris[, c("Species", "Sepal.Length", "Sepal.Width")]
#' d$Species <- as.character(d$Species)
#' d$Species[d$Species != "setosa"] <- "non-setosa"
#' x <- d$Sepal.Length
#' x[d$Sepal.Length <= 5.2] <- "Very Short"
#' x[d$Sepal.Length >  5.2 & d$Sepal.Length <= 6.1] <- "Short"
#' x[d$Sepal.Length >  6.1 & d$Sepal.Length <= 7.0] <- "Long"
#' x[d$Sepal.Length >  7.0] <- "Very Long"
#' d$Sepal.Length <- x
#' decTree <- decisionTree(d, eta=5, purity=0.95, minsplit=0)
#' plot(decTree)

plot.DecisionTreeObject <- function(DecisionTreeObject, dict, green=1, left=1) {
  # ---- prepare data
  treeResult <- DecisionTreeObject@resultDF
  nLevel <- max(treeResult$level)
  d <- cbind(treeResult, x=0, y=0)
  d$x[d$leaf != 0] <- d$leaf[d$leaf != 0]
  d$x[d$leaf != 0] <- d$leaf[d$leaf != 0]
  d$y <- d$level

  # ---- set values on leaves
  d <- cbind(d, name=as.character(d$vName), stringsAsFactors=F)
  if (left == 1) {
    d$name[d$leaf != 0] <- paste0(d$Lleft[d$leaf != 0], "/", 
                                  d$Lsize[d$leaf != 0], "\n",
                                  round(d$Lleft[d$leaf != 0] / 
                                        d$Lsize[d$leaf != 0] * 100), "%")
  } else {
    d$name[d$leaf != 0] <- paste0(d$Lright[d$leaf != 0], "/", 
                                  d$Lsize[d$leaf != 0], "\n",
                                  round(d$Lright[d$leaf != 0] / 
                                        d$Lsize[d$leaf != 0] * 100), "%")

  }

  for (i in (nLevel - 1):1) {
    for (j in d$node) {
      if (d$level[j] == i & d$leaf[j] == 0) {
        crow <- d[d$node == j,]
        d$x[j] <- (d$x[d$parent == j & d$LR == "L"] + 
                   d$x[d$parent == j & d$LR == "R"]) / 2
      }
    }
  }

  # ---- plot edges
  plot.new()
  plot.window(c(0.5, nrow(d[d$leaf != 0,]) + 0.5), c(nLevel, 1))

  d <- cbind(d, parx=0, pary=0)
  for (i in 2:nrow(d)) {
    cpar <- d[d$node == d[i, "parent"],]  # parent
    width <- (d$Lsize[i] + d$Rsize[i]) / (d[1,"Lsize"] + d[1, "Rsize"]) * 80
    y <- seq(0, pi, length.out=30)
    x <- cos(y)
    lines((x / 2 + 1/2) * (cpar$x - d[i, "x"]) + d[i, "x"],
          y / pi + cpar$level, lwd = width, col="gray") 
  }

  # ---- plot vertices
  rect(0.5, 0, nrow(d[d$leaf != 0,]) + 0.5, 1, col="white", border=NA)
  color <- rep("white", nrow(d))
  cols <- (d$Lleft + d$Rleft) / (d$Lsize + d$Rsize) * 255 
  if (green == 1) {
    color <- paste0("#", format(as.hexmode(round(255 - cols)), 2), 
                    format(as.hexmode(round(cols)), 2), "00")
  } else {
    color <- paste0("#", format(as.hexmode(round(cols)), 2), 
                    format(as.hexmode(round(255 - cols)), 2), "00")
  }
  widthCoef <- rep(1.8, nrow(d))
  widthCoef[d$leaf != 0] <- 1.1
  if (missing(dict)) {
    widths <- strwidth(d$name)/widthCoef
    widths[widths < 0.45 & d$leaf == 0] <- 0.45
    rect(d$x - widths, d$y - 0.15 * (nLevel - 1) / 3,  
         d$x + widths, d$y + 0.15 * (nLevel - 1) / 3, 
         col=color, border="gray")
    text(d$x, d$y, d$name)
  } else {
    dd <- merge(d, dict, by="vName", all=T)
    dd$toPlot[dd$vName == "leaf"] <- dd$name[dd$vName == "leaf"]
    dd <- dd[order(dd$node),]
    widths <- strwidth(dd$toPlot)/widthCoef
    widths[widths < 0.45 & d$leaf == 0] <- 0.45
    rect(dd$x - widths, dd$y - 0.15 * (nLevel - 1) / 3,
         dd$x + widths, dd$y + 0.15 * (nLevel - 1) / 3, 
         col=color, border="gray")
    text(dd$x, dd$y, dd$toPlot)
  }

  # ---- plot conditions on edges
  for (i in 1:nrow(d)) {
    if (d$v[i] != "") {
      Lchild <- d[d$parent == i & d$LR == "L",]
      Rchild <- d[d$parent == i & d$LR == "R",]
      x <- d$x[i] + (Lchild$x - d$x[i] ) / 2 - 0.1  # TODO number of leaves
      y <- d$y[i] + (Lchild$y - d$y[i] ) / 2 
      text(x, y, d$v[i]) 
      x <- d$x[i] + (Rchild$x - d$x[i] ) / 2 + 0.1
      y <- d$y[i] + (Rchild$y - d$y[i] ) / 2 
      text(x, y, d$vC[i])
    }
  }
}