Allow a function as lengths in tree-generation functions

NEWS.md

@@ -2,6 +2,11 @@
 
 ## New features
 
+- `RandomTree()`, `YuleTree()`, `PectinateTree()`, `BalancedTree()`,
+  `StarTree()`, and `SingleTaxonTree()` accept a function as `lengths`;
+  the function is called with the number of edges as its sole argument
+  (e.g. `RandomTree(8, lengths = runif)`).
+
 - `NexusTokensToInteger()` converts character data to integers,
   mapping uncertain tokens to `NA`.
 - `PaintTree()` assigns colours to every edge, leaf, and internal node such

R/phylo.R

@@ -64,7 +64,8 @@ Renumber <- function(tree) {
 #' `ZeroTaxonTree()` creates an empty `phylo` object with zero leaves or edges.
 #'
 #' @param label a character vector specifying the label of the tip.
-#' @param lengths a numeric vector specifying the edge lengths of the tree.
+#' @param lengths Numeric vector of edge lengths, or a function called with the
+#' number of edges as its argument (e.g. `lengths = runif`).
 #' @return `SingleTaxonTree()` returns a \code{phylo} object containing a single
 #' tip with the specified label.
 #'
@@ -81,15 +82,12 @@ NULL
 #' @rdname TrivialTree
 #' @export
 SingleTaxonTree <- function(label = "t1", lengths = NULL) {
-  if (is.null(lengths)) {
-    structure(list(edge = matrix(c(2L, 1L), 1, 2), tip.label = label,
-                   Nnode = 1L),
-              class = "phylo", order = "preorder")
-  } else {
-    structure(list(edge = matrix(c(2L, 1L), 1, 2), tip.label = label,
-                   Nnode = 1L, edge.length = lengths[[1]]),
-              class = "phylo", order = "preorder")
+  tr <- list(edge = matrix(c(2L, 1L), 1, 2), tip.label = label, Nnode = 1L)
+  if (!is.null(lengths)) {
+    tr[["edge.length"]] <- .EdgeLengths(lengths, 1L)[[1L]]
   }
+  # Return:
+  structure(tr, class = "phylo", order = "preorder")
 }
 
 #' @rdname TrivialTree

R/tree_generation.R

@@ -18,6 +18,10 @@
 #' @name GenerateTree
 NULL
 
+.EdgeLengths <- function(lengths, nEdge) {
+  if (is.function(lengths)) lengths(nEdge) else rep(lengths, length.out = nEdge)
+}
+
 #' @rdname GenerateTree
 #'
 #' @param root Character or integer specifying tip to use as root;
@@ -41,6 +45,8 @@ NULL
 #' data("Lobo")
 #' RandomTree(Lobo.phy)
 #'
+#' RandomTree(8, lengths = runif)
+#'
 #' @export
 RandomTree <- function(tips, root = FALSE, nodes, lengths = NULL) {
   tips <- TipLabels(tips)
@@ -107,9 +113,9 @@ RandomTree <- function(tips, root = FALSE, nodes, lengths = NULL) {
   }
 
   if (!is.null(lengths)) {
-    tree[["edge.length"]] <- rep(lengths, length.out = dim(tree[["edge"]])[[1]])
+    tree[["edge.length"]] <- .EdgeLengths(lengths, dim(tree[["edge"]])[[1]])
   }
-  
+
   # Return:
   tree
 }
@@ -153,9 +159,9 @@ YuleTree <- function(tips, addInTurn = FALSE, root = TRUE, lengths = NULL) {
   }
 
   if (!is.null(lengths)) {
-    tree[["edge.length"]] <- rep(lengths, length.out = dim(tree[["edge"]])[[1]])
+    tree[["edge.length"]] <- .EdgeLengths(lengths, dim(tree[["edge"]])[[1]])
   }
-  
+
   # Return:
   tree
 }
@@ -211,7 +217,7 @@ PectinateTree <- function(tips, lengths = NULL) {
         tip.label = tips
       )
     if (!is.null(lengths) && nTip > 1) {
-      tr[["edge.length"]] <- rep(lengths, length.out = 2 * (nTip - 1))
+      tr[["edge.length"]] <- .EdgeLengths(lengths, 2L * (nTip - 1L))
     }
     structure(tr, order = "cladewise", class = "phylo")
   }
@@ -242,7 +248,7 @@ BalancedTree <- function(tips, lengths = NULL) {
                Nnode = nTip - 1L,
                tip.label = as.character(tips))
     if (!is.null(lengths)) {
-      tr[["edge.length"]] <- rep(lengths, length.out = 2 * (nTip - 1))
+      tr[["edge.length"]] <- .EdgeLengths(lengths, 2L * (nTip - 1L))
     }
     # Return:
     structure(tr, order = "preorder", class = "phylo")
@@ -284,19 +290,13 @@ StarTree <- function(tips, lengths = NULL) {
   parent <- rep.int(nTip + 1L, nTip)
   child <- seq_len(nTip)
 
-  tr <- if (is.null(lengths)) {
-    list(
-      edge = matrix(c(parent, child), ncol = 2L),
-      Nnode = 1L,
-      tip.label = tips
-    )
-  } else {
-    list(
-      edge = matrix(c(parent, child), ncol = 2L),
-      Nnode = 1L,
-      tip.label = tips,
-      edge.length =  rep(lengths, length.out = nTip)
-    )
+  tr <- list(
+    edge = matrix(c(parent, child), ncol = 2L),
+    Nnode = 1L,
+    tip.label = tips
+  )
+  if (!is.null(lengths)) {
+    tr[["edge.length"]] <- .EdgeLengths(lengths, nTip)
   }
   structure(tr, order = "cladewise", class = "phylo")
 }

tests/testthat/test-tree_generation.R

@@ -38,12 +38,14 @@ test_that("BalancedTree(lengths)", {
   expect_equal(BalancedTree(1, lengths = 2), SingleTaxonTree("t1", lengths = 2))
   expect_equal(BalancedTree(2, lengths = 2)$edge.length, c(2, 2))
   expect_equal(BalancedTree(3, lengths = 1:3)$edge.length, c(1:3, 1))
+  expect_equal(BalancedTree(3, lengths = seq_len)$edge.length, 1:4)
 })
 test_that("PectinateTree(lengths)", {
   expect_equal(PectinateTree(0, lengths = 1), ZeroTaxonTree())
   expect_equal(PectinateTree(1, lengths = 2), SingleTaxonTree("t1", lengths = 2))
   expect_equal(PectinateTree(2, lengths = 2)$edge.length, c(2, 2))
   expect_equal(PectinateTree(3, lengths = 1:3)$edge.length, c(1:3, 1))
+  expect_equal(PectinateTree(3, lengths = seq_len)$edge.length, 1:4)
 })
 
 test_that("StarTree() works", {
@@ -52,6 +54,7 @@ test_that("StarTree() works", {
   expect_true(is.integer(StarTree(8)$edge))
   expect_null(StarTree(8L)[["edge.length"]])
   expect_equal(StarTree(8L, 8:1)[["edge.length"]], 8:1)
+  expect_equal(StarTree(4L, seq_len)[["edge.length"]], 1:4)
 })
 
 test_that("Random trees are generated correctly", {
@@ -100,6 +103,8 @@ test_that("RandomTree(lengths)", {
   expect_equal(RandomTree(3, lengths = 1:3, root = FALSE)$edge.length, 1:3)
   expect_equal(RandomTree(3, lengths = 1:3, root = TRUE)$edge.length, c(1:3, 1))
   expect_equal(RandomTree(5, nodes = 2, lengths = 1:3)$edge.length, c(1:3, 1:3))
+  expect_equal(RandomTree(3, lengths = seq_len, root = FALSE)$edge.length, 1:3)
+  expect_equal(RandomTree(3, lengths = seq_len, root = TRUE)$edge.length, 1:4)
 })
 
 test_that("Small random trees are generated", {
@@ -147,6 +152,12 @@ test_that("YuleTree(lengths)", {
   expect_equal(YuleTree(3, lengths = 2)$edge.length, rep(2, 4))
   expect_equal(YuleTree(3, lengths = 1:3, root = FALSE)$edge.length, 1:3)
   expect_equal(YuleTree(3, lengths = 1:3, root = TRUE)$edge.length, c(1:3, 1))
+  expect_equal(YuleTree(3, lengths = seq_len, addInTurn = TRUE)$edge.length, 1:4)
+})
+
+test_that("SingleTaxonTree(lengths)", {
+  expect_equal(SingleTaxonTree("t1", lengths = 2)$edge.length, 2)
+  expect_equal(SingleTaxonTree("t1", lengths = seq_len)$edge.length, 1)
 })
 
 test_that("Hamming() works", {