R/dendro_tree.R

In ggdendro: Create Dendrograms and Tree Diagrams Using 'ggplot2'

#
#  ggdendro/R/dendro_tree.R by Andrie de Vries  Copyright (C) 2011-2015
#  Contains code from tree/R/tree.R by B. D. Ripley  Copyright (C) 1994-2015
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 or 3 of the License
#  (at your option).
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  A copy of the GNU General Public License is available at
#  http://www.r-project.org/Licenses/
#
#


# tree_ggdendro_env <- new.env()

#' Extract data from regression tree object for plotting using ggplot.
#'
#' Extracts data to plot line segments and labels from a `tree::tree()` object.
#' This data can then be manipulated or plotted, e.g. using [ggplot2::ggplot()].
#'
#' @param model object of class "tree", e.g. the output of tree()
#'
#' @param type Either `proportional` or `uniform`. If this partially matches
#'   "uniform", the branches are of uniform length. Otherwise they are
#'   proportional to the decrease in impurity.
#'
#' @param ... ignored
#'
#' @export
#'
#' @return
#' A list of three data frames:
#' \item{segments}{a data frame containing the line segment data}
#' \item{labels}{a data frame containing the label text data}
#' \item{leaf_labels}{a data frame containing the leaf label text data}
#'
#' @seealso [ggdendrogram()]
#' @family dendro_data methods
#' @family tree functions
#' @author Andrie de Vries, using code modified from original by Brian Ripley
#' @example inst/examples/example_dendro_tree.R
dendro_data.tree <- function(model, type = c("proportional", "uniform"), ...) {
  type <- match.arg(type)
  uniform <- type == "uniform"

  # dev <- dev.cur()
  # if (dev == 1L) dev <- 2L # as device will be opened.

  # assign(paste0("device", dev), uniform, envir = tree_ggdendro_env)

  labels <- tree_labels(model, uniform = uniform, ...)
  as.dendro(
    segments = tree_segments(model, uniform, ...),
    labels = labels$labels,
    leaf_labels = labels$leaf_labels,
    class = "tree"
  )
}



#' Extract data frame from tree object for plotting using ggplot.
#' @param model object of class "tree", e.g. the output of tree()
#' @param ... ignored
#' @keywords internal
#' @seealso [ggdendrogram()]
#' @family tree functions
#' @author Code modified from original by Brian Ripley
tree_segments <- function(model, uniform, ...) {
  if (missing(uniform)) stop("specify the uniform argument")
  # Uses tree:::treeco to extract data frame of plot locations
  xy <- treeco(model, uniform = uniform)
  n <- model$frame$n

  # Lines copied from tree:::treepl
  x <- xy$x
  y <- xy$y
  node <- as.numeric(row.names(model$frame))
  parent <- match((node %/% 2), node)
  sibling <- match(ifelse(node %% 2, node - 1L, node + 1L), node)

  linev <- data.frame(x = x, y = y, xend = x, yend = y[parent], n = n)
  lineh <- data.frame(x = x[parent], y = y[parent], xend = x, yend = y[parent], 
                      n = n)

  rbind(linev[-1, ], lineh[-1, ])
}

#' Extract labels data frame from tree object for plotting using ggplot.
#' @param model object of class "tree", e.g. the output of tree()
#' @param ... ignored
#' @return a list with two elements: $labels and $leaf_labels
#' @keywords internal
#' @seealso [ggdendrogram()]
#' @family tree functions
#' @author Code modified from original by Brian Ripley
tree_labels <- function(model, uniform, ...) {
  # Uses tree:::treeco to extract data frame of plot locations
  xy <- treeco(model, uniform = uniform)
  label <- model$frame$var
  yval <- model$frame$yval
  sleft <- model$frame$splits.cutleft
  sright <- model$frame$splits.right

  # Lines copied from tree:::treepl
  x <- xy$x
  y <- xy$y
  node <- as.numeric(row.names(model$frame))
  parent <- match((node %/% 2), node)
  sibling <- match(ifelse(node %% 2, node - 1L, node + 1L), node)

  # Extract labels
  data <- data.frame(x = x, y = y, label = label)
  data <- data[data$label != "<leaf>", ]
  labels <- as.data.frame(data)

  # Extract leaf labels
  data <- data.frame(x, y, label, yval)
  data <- data[data$label == "<leaf>", ]
  if (is.numeric(data$yval)) {
    data$label <- round(data$yval, 2)
  } else {
    data$label <- data$yval
  }
  leaf_labels <- as.data.frame(data)

  list(
    labels = labels,
    leaf_labels = leaf_labels
  )
}

#' Extract labels data frame from tree object for plotting using ggplot.
#'
#' Extract labels data frame from tree object for plotting using ggplot
#'
#' @param model object of class "tree", e.g. the output of tree()
#' @param ... ignored
#' @keywords internal
#' @seealso [ggdendrogram()]
#' @family tree functions
#' @author Code modified from original by Brian Ripley
get_data_tree_leaf_labels <- function(model, uniform, ...) {
  # Uses tree:::treeco to extract data frame of plot locations
  xy <- treeco(model, uniform = uniform)
  label <- model$frame$var
  yval <- model$frame$yval
  sleft <- model$frame$splits.cutleft
  sright <- model$frame$splits.right

  # Lines copied from tree:::treepl
  x <- xy$x
  y <- xy$y
  node <- as.numeric(row.names(model$frame))
  parent <- match((node %/% 2), node)
  sibling <- match(ifelse(node %% 2, node - 1L, node + 1L), node)

  data <- data.frame(x, y, label, yval)
  data <- data[data$label == "<leaf>", ]
  data$label <- round(data$yval, 2)
  data
}


# treeco ------------------------------------------------------------------

#' Function copied from tree:::treeco.
#'
#' @param tree tree object
#' @param uniform ???
#' @keywords internal
treeco <- function(tree, uniform) {
  # if (missing(uniform)) {
  #   pn <- paste0("device", dev.cur())
  #   uniform <- if (exists(pn, envir = tree_ggdendro_env, inherits = FALSE))
  #     get(pn, envir = tree_ggdendro_env, inherits = FALSE)
  #   else FALSE
  # }
  if (missing(uniform)) stop("specify uniform argument")

  frame <- tree$frame
  node <- as.integer(row.names(frame))
  depth <- tree.depth(node)
  x <- -depth
  if (uniform) {
    y <- x
  } else {
    y <- dev <- frame$dev
    depth <- split(seq(node), depth)
    parent <- match(node %/% 2L, node)
    sibling <- match(
      ifelse(node %% 2L, node - 1L, node + 1L),
      node
    )
    for (i in depth[-1L]) {
      y[i] <- y[parent[i]] - dev[parent[i]] +
        dev[i] + dev[sibling[i]]
    }
  }
  depth <- -x
  leaves <- frame$var == "<leaf>"
  x[leaves] <- seq(sum(leaves))
  depth <- split(seq(node)[!leaves], depth[!leaves])
  left.child <- match(node * 2L, node)
  right.child <- match(node * 2 + 1L, node)
  for (i in rev(depth)) x[i] <- 0.5 * (x[left.child[i]] + x[right.child[i]])
  list(x = x, y = y)
}