R/get_nodes_xy.R

In talgalili/dendextend: Extending 'dendrogram' Functionality in R

# Copyright (C) Tal Galili
#
# This file is part of dendextend.
#
# dendextend 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 of the License, or
# (at your option) any later version.
#
# dendextend 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/
#




# We had to write a stripped down version of the following functions:
# stats:::plot.dendrogram
# stats:::plotNodeLimit
# stats:::plotNode
# stats:::.memberDend(dend)
# stats:::.midDend(dend)




#' @title Get the x-y coordinates of a dendrogram's nodes
#' @description Get the x-y coordinates of a dendrogram's nodes. Can be used to add text or images on the tree.
#' @export
#' @param dend a dendrogram object
#' @param type type of plot.
#' @param center logical; if TRUE, nodes are plotted centered with respect to the
#' leaves in the branch. Otherwise (default),
#' plot them in the middle of all direct child nodes.
#' @param horiz logical indicating if the dendrogram should be drawn horizontally or not.
#' @param ... not used
#' @source This is a striped down version of the
#' function \code{\link{plot.dendrogram}}.
#' It performs (almost) the same task, only it does not do any plotting
#' but it does save the x-y coordiantes of the nodes.
#' @return
#' A 2-dimensional matrix, with rows as the number of nodes,
#' and the first column is the x location, while the second is the
#' y location.
#' @seealso \link{get_nodes_attr}, \link{nnodes},
#' \link{nleaves}
#' @examples
#' \dontrun{
#'
#' # If we would like to see the numbers from plot:
#' # ?getOption("verbose")
#' # options(verbose=TRUE)
#' # options(verbose=FALSE)
#'
#' # -----
#' # Draw a depth first search illustration
#' # -----
#'
#' dend <- 1:5 %>%
#'   dist() %>%
#'   hclust() %>%
#'   as.dendrogram()
#' get_nodes_xy(dend)
#'
#' # polygon(get_nodes_xy(dend), col = 2)
#' plot(dend,
#'   leaflab = "none",
#'   main = "Depth-first search in a dendrogram"
#' )
#' xy <- get_nodes_xy(dend)
#' for (i in 1:(nrow(xy) - 1)) {
#'   arrows(xy[i, 1], xy[i, 2],
#'     angle = 17,
#'     length = .5,
#'     xy[i + 1, 1], xy[i + 1, 2],
#'     lty = 1, col = 3, lwd = 1.5
#'   )
#' }
#' points(xy, pch = 19, cex = 4)
#' text(xy, labels = 1:nnodes(dend), cex = 1.2, col = "white", adj = c(0.4, 0.4))
#' }
get_nodes_xy <- function(dend, type = c("rectangle", "triangle"), center = FALSE,
                         horiz = FALSE, ...) {
  x <- dend # TODO: this code could be cleaned

  # this function/env let's us keep a "stack" of the xy matrix, and update it
  # no matter how deep we get in the recursion...
  nodes_xy <- local({

    # notice how we must create this function within get_nodes_xy
    # since it relies on the specific value of "x" when first created.
    xy_matrix <- matrix(0, nrow = nnodes(x), ncol = 2)
    i_node <- 0

    function(xy) {
      if (missing(xy)) {
        return(with(environment(nodes_xy), xy_matrix))
      }

      # next node
      i_node <<- i_node + 1
      xy_matrix[i_node, ] <<- xy
    }
  })


  type <- match.arg(type)
  hgt <- attr(x, "height")
  mem.x <- nleaves(x)

  if (center) {
    x1 <- 0.5
    x2 <- mem.x + 0.5
  }
  else {
    x1 <- 1
    x2 <- mem.x
  }

  plotNode2(x1, x2, x,
    type = type, center = center, leaflab = "none",
    horiz = horiz,
    nodes_xy = nodes_xy
  )

  return(nodes_xy())
}


















plotNodeLimit2 <- function(x1, x2, subtree, center) {
  inner <- !is.leaf(subtree) && x1 != x2
  if (inner) {
    K <- length(subtree)
    mTop <- nleaves(subtree)
    limit <- integer(K)
    xx1 <- x1
    for (k in 1L:K) {
      m <- nleaves(subtree[[k]])
      xx1 <- xx1 + (if (center) {
        (x2 - x1) * m / mTop
      } else {
        m
      })
      limit[k] <- xx1
    }
    limit <- c(x1, limit)
  }
  else {
    limit <- c(x1, x2)
  }
  mid <- attr(subtree, "midpoint")
  center <- center || (inner && !is.numeric(mid))
  x <- if (center) {
    mean(c(x1, x2))
  } else {
    x1 + (if (inner) {
      mid
    } else {
      0
    })
  }
  list(x = x, limit = limit)
}



.midDend2 <- function(x)
  if (is.null(mp <- attr(x, "midpoint"))) 0 else mp


plotNode2 <- function(x1, x2, subtree, type, center, leaflab, horiz = FALSE, nodes_xy) {
  inner <- !is.leaf(subtree) && x1 != x2
  yTop <- attr(subtree, "height")
  bx <- plotNodeLimit2(x1, x2, subtree, center)
  xTop <- bx$x

  if (getOption("verbose")) {
    cat(if (inner) {
      "inner node"
    } else {
      "leaf"
    }, ":")
    cat(if (inner) {
      paste(" height", formatC(yTop), "; ")
    }, "(x1,x2)= (",
    formatC(x1, width = 4), ",", formatC(x2, width = 4),
    ")", "--> xTop=", formatC(xTop, width = 8), "\n",
    sep = ""
    )
  }

  # update the xy coordinates using a function defined in the beginning!
  nodes_xy(c(xTop, yTop))



  i <- if (inner) {
    1
  } else {
    2
  }

  if (inner) {
    for (k in seq_along(subtree)) {
      child <- subtree[[k]]
      yBot <- attr(child, "height")
      if (getOption("verbose")) {
        cat("ch.", k, "@ h=", yBot, "; ")
      }
      if (is.null(yBot)) {
        yBot <- 0
      }
      xBot <- if (center) {
        mean(bx$limit[k:(k + 1)])
      } else {
        bx$limit[k] + .midDend2(child)
      }

      i <- if (!is.leaf(child)) {
        1
      } else {
        2
      }

      vln <- NULL

      if (!is.null(attr(child, "edgetext"))) {
        if (getOption("verbose")) {
          #                cat("-- with \"edgetext\"", format(edgeText))
          if (!is.null(vln)) {
            mx <- if (type == "triangle") {
              (xTop + xBot + ((xTop - xBot) / (yTop - yBot)) *
                vln) / 2
            } else {
              xBot
            }
            my <- (yTop + yBot + 2 * vln) / 2
          }
          else {
            mx <- if (type == "triangle") {
              (xTop + xBot) / 2
            } else {
              xBot
            }
            my <- (yTop + yBot) / 2
          }
        }

        #             vlm <- strheight(c(edgeText, "h"), cex = t.cex)/2
        #             hlm <- strwidth(c(edgeText, "m"), cex = t.cex)/2
        #             hl3 <- c(hlm[1L], hlm[1L] + hlm[2L], hlm[1L])
      }
      plotNode2(bx$limit[k], bx$limit[k + 1],
        subtree = child,
        type, center, leaflab, horiz, nodes_xy = nodes_xy
      )
    }
  }
  invisible()
}