R/heat_tree.R

In treeheatr: Heatmap-Integrated Decision Tree Visualizations

#' Draws and aligns decision tree and heatmap.
#'
#' @param x Dataframe or a `party` or `partynode` object representing a custom tree.
#' If a dataframe is supplied, conditional inference tree is computed.
#' If a custom tree is supplied, it must follow the partykit syntax:
#' https://cran.r-project.org/web/packages/partykit/vignettes/partykit.pdf
#' @param data_test Tidy test dataset. Required if `x` is a `partynode` object.
#' If NULL, heatmap displays (training) data `x`.
#' @param target_lab Name of the column in data that contains target/label information.
#' @param task Character string indicating the type of problem,
#' either 'classification' (categorical outcome) or 'regression' (continuous outcome).
#' @param feat_types Named vector indicating the type of each features,
#' e.g., c(sex = 'factor', age = 'numeric').
#' If feature types are not supplied, infer from column type.
#' @param label_map Named vector of the meaning of the target values,
#' e.g., c(`0` = 'Edible', `1` = 'Poisonous').
#' @param target_cols Character vectors representing the hex values of different
#' level colors for targets, defaults to viridis option B.
#' @param target_legend Logical. If TRUE, target legend is drawn.
#' @param clust_samps Logical. If TRUE, hierarchical clustering would be performed
#' among samples within each leaf node.
#' @param clust_target Logical. If TRUE, target/label is included in hierarchical clustering
#' of samples within each leaf node and might yield a more interpretable heatmap.
#' @param custom_layout Dataframe with 3 columns: id, x and y
#' for manually input custom layout.
#' @param show Character string indicating which components of the decision tree-heatmap
#' should be drawn. Can be 'heat-tree', 'heat-only' or 'tree-only'.
#' @param heat_rel_height Relative height of heatmap compared to whole figure (with tree).
#' @param lev_fac Relative weight of child node positions
#' according to their levels, commonly ranges from 1 to 1.5.
#' 1 for parent node perfectly in the middle of child nodes.
#' @param panel_space Spacing between facets relative to viewport,
#' recommended to range from 0.001 to 0.01.
#' @param print_eval Logical. If TRUE, print evaluation of the tree performance.
#' Defaults to TRUE when `data_test` is supplied.
#' @param \dots Further arguments passed to `draw_tree()` and/or `draw_heat()`.
#'
#' @return A gtable/grob object of the decision tree (top) and heatmap (bottom).
#' @export
#'
#' @examples
#'
#' \donttest{
#' heat_tree(penguins, target_lab = "species")
#' heat_tree(
#'   x = galaxy[1:100, ],
#'   target_lab = "target",
#'   task = "regression",
#'   terminal_vars = NULL,
#'   tree_space_bottom = 0
#' )
#' }
heat_tree <- function(
  x,
  target_lab = NULL,
  data_test = NULL,
  task = c("classification", "regression"),
  feat_types = NULL,
  label_map = NULL,
  target_cols = NULL,
  target_legend = FALSE,
  clust_samps = TRUE,
  clust_target = TRUE,
  custom_layout = NULL,
  show = "heat-tree",
  heat_rel_height = 0.2,
  lev_fac = 1.3,
  panel_space = 0.001,
  print_eval = (!is.null(data_test)),
  ...
) {
  target_cols <- get_cols(
    target_cols,
    match.arg(task),
    guide = ifelse(target_legend, "legend", "none")
  )

  mf <- match.call()
  ctree_vars <-
    c(
      "x",
      "data_test",
      "target_lab",
      "task",
      "feat_types",
      "label_map",
      "clust_samps",
      "clust_target",
      "custom_layout",
      "lev_fac",
      "panel_space"
    )
  tree_vars <-
    c(
      "title",
      "task",
      "tree_space_top",
      "tree_space_bottom",
      "par_node_vars",
      "terminal_vars",
      "edge_vars",
      "edge_text_vars",
      "print_eval",
      "metrics",
      "x_eval",
      "y_eval"
    )
  heat_vars <-
    c(
      "feat_types",
      "trans_type",
      "cont_cols",
      "cate_cols",
      "clust_feats",
      "cont_legend",
      "cate_legend",
      "target_space",
      "panel_space",
      "target_pos",
      "feats",
      "show_all_feats",
      "p_thres",
      "target_lab_disp"
    )

  m_ctree <- match(ctree_vars, names(mf), 0L)
  m_tree <- match(tree_vars, names(mf), 0L)
  m_heat <- match(heat_vars, names(mf), 0L)

  ################################################################
  ##### Compute conditional inference tree:

  ctree_result <- mf[c(1L, m_ctree)]
  ctree_result[[1L]] <- quote(compute_tree)
  ctree_result <- eval(ctree_result, parent.frame())

  ################################################################
  ##### Draw decision tree and heatmap:

  dtree <- mf[c(1L, m_tree)]
  dtree$target_cols <- target_cols
  dtree$print_eval <- print_eval
  for (argg in c("fit", "dat", "layout", "term_dat")) {
    dtree[[argg]] <- ctree_result[[argg]]
  }
  dtree[[1L]] <- quote(draw_tree)
  dtree <- eval(dtree, parent.frame())

  dheat <- mf[c(1L, m_heat)]
  dheat$fit <- ctree_result$fit
  dheat$dat <- ctree_result$dat
  dheat$target_cols <- target_cols
  dheat[[1L]] <- quote(draw_heat)
  dheat <- eval(dheat, parent.frame())

  ################################################################
  ##### Align decision tree and heatmap:

  g <- align_plots(dheat, dtree, heat_rel_height, show)
  class(g) <- c("ggHeatTree", class(g))
  g
}