Node numbers in CollapseNode() (#151)

NEWS.md

@@ -4,10 +4,10 @@
 - `RoguePlot()$legendLabels` returns suggested labels for legend.
 - `DescendantTips()` complements `DescendantEdges()`.
 - `NodeNumbers()` returns the indices of nodes within a tree.
-- `AddTip()` supports node labels.
+- Support node labels in `AddTip()`, `CollapseNode()`
+  ([#149](https://github.com/ms609/TreeTools/issues/149)).
 - `AddTip(edgeLength = NULL)` defaults to `lengthBelow`. This will become the
   default in a future release.
-  ([#149](https://github.com/ms609/TreeTools/issues/149)).
 - An entry point to the C++ function `root_on_node()` is now exported
   (intended for expert use only).
 - Use `KeepTip()` internally so `SplitFrequency()` supports `Splits` objects

R/tree_rearrangement.R

@@ -403,6 +403,8 @@ CollapseNode.phylo <- function(tree, nodes) {
   edge <- tree[["edge"]]
   lengths <- tree[["edge.length"]]
   hasLengths <- !is.null(lengths)
+  nodeLabels <- tree[["node.label"]]
+  hasLabels <- !is.null(nodeLabels)
   parent <- edge[, 1]
   child <- edge[, 2]
   root <- RootNode(edge)
@@ -428,7 +430,9 @@ CollapseNode.phylo <- function(tree, nodes) {
 
   for (node in nodes[order(depths)]) {
     newParent <- parent[edgeBelow[node]]
-    if (hasLengths) lengths[parent == node] <- lengths[parent == node] + lengths[child == node]
+    if (hasLengths) {
+      lengths[parent == node] <- lengths[parent == node] + lengths[child == node]
+    }
     parent[parent == node] <- newParent
   }
 
@@ -438,6 +442,9 @@ CollapseNode.phylo <- function(tree, nodes) {
                           newNumber[child[keptEdges]])
   tree[["edge.length"]] <- lengths[keptEdges]
   tree[["Nnode"]] <- tree[["Nnode"]] - length(nodes)
+  if (hasLabels) {
+    tree[["node.label"]] <- nodeLabels[-(nodes - nTip)]
+  }
 
   # TODO Renumber nodes sequentially
   # TODO Re-write this in C++.

tests/testthat/test-tree_rearrange.R

@@ -296,6 +296,42 @@ test_that("CollapseNode() works", {
   expect_error(CollapseEdge(tree, 9))
 })
 
+test_that("CollapseNode() handles node labels", {
+  bal6 <- BalancedTree(6)
+  startLabels <- paste("Node", 7:11)
+  bal6[["node.label"]] <- startLabels
+  if (interactive()) {
+    plot(bal6, show.node.label = TRUE)
+  }
+
+  # Collapse a cherry
+  expect_equal(
+    CollapseNode(bal6, 9)[["node.label"]],
+    startLabels[-3]
+  )
+
+  # Collapse an internal node
+  expect_equal(
+    CollapseNode(bal6, 8)[["node.label"]],
+    startLabels[-2]
+  )
+
+  # Collapse an internal node
+  expect_equal(
+    CollapseEdge(bal6, 1)[["node.label"]],
+    startLabels[-2]
+  )
+
+  # case = 3 -> y is bound on an internal edge
+  expect_equal(
+    CollapseNode(bal6, c(8, 11))[["node.label"]],
+    startLabels[-c(2, 5)]
+  )
+
+  expect_equal(AddTipEverywhere(bal6)[[1]][["node.label"]],
+               AddTip(bal6, where = 1)[["node.label"]])
+})
+
 test_that("Binarification is uniform", {
   set.seed(0)
   Test <- function(tree, nTree, nSamples = 200L, ape = FALSE) {