Optimization read.mcmctree

NAMESPACE

@@ -22,6 +22,8 @@ S3method(as.treedata,phylo4d)
 S3method(as.treedata,pml)
 S3method(as.treedata,pvclust)
 S3method(as.treedata,tbl_df)
+S3method(as.ultrametric,phylo)
+S3method(as.ultrametric,treedata)
 S3method(as_tibble,pvclust)
 S3method(child,phylo)
 S3method(child,treedata)
@@ -59,6 +61,7 @@ export(Ntip)
 export(ancestor)
 export(as.phylo)
 export(as.treedata)
+export(as.ultrametric)
 export(as_tibble)
 export(child)
 export(drop.tip)
@@ -131,9 +134,11 @@ importFrom(ape,as.AAbin.character)
 importFrom(ape,as.DNAbin)
 importFrom(ape,as.DNAbin.character)
 importFrom(ape,as.phylo)
+importFrom(ape,branching.times)
 importFrom(ape,checkLabel)
 importFrom(ape,extract.clade)
 importFrom(ape,is.rooted)
+importFrom(ape,is.ultrametric)
 importFrom(ape,read.FASTA)
 importFrom(ape,read.nexus)
 importFrom(ape,read.tree)

R/MCMCTree.R

@@ -2,13 +2,17 @@
 ##'
 ##' @title read.mcmctree
 ##' @param file the output tree file of MCMCTree
+##' @param force.ultrametric logical whether convert the tree 
+##' to be ultrametric, if it is not ultrametric, default is FALSE.
+##' When the tree is ultrametric, branch times will be calculated 
+##' automatically. 
 ##' @return treedata object
 ##' @export
 ##' @examples
 ##' file <- system.file("extdata/MCMCTree", "mcmctree_output.tree", package="treeio")
 ##' tr <- read.mcmctree(file)
 ##' tr
-read.mcmctree <- function(file){
+read.mcmctree <- function(file, force.ultrametric = FALSE){
     text <- readLines(file)
     ind <- grep("^.*tree.*=.*", text, ignore.case=TRUE)
     text[ind] <- gsub("^.*TREE", "TREE", text[ind], ignore.case=TRUE)
@@ -18,10 +22,36 @@ read.mcmctree <- function(file){
     obj <- read.beast(file=newfile)
     if(inherits(obj, "treedata")){
         obj@file <- filename(file)
-    } else{
+        obj <- add.branch.time.mcmctree(
+                                        obj=obj, 
+                                        force.ultrametric = force.ultrametric, 
+                                        )
+    }else{
         for (i in seq_len(length(obj))){
             obj[[i]]@file <- filename(file)
+            obj[[i]] <- add.branch.time.mcmctree(
+                                        obj=obj[[i]], 
+                                        force.ultrametric = force.ultrametric
+                                        )
         }
     }
     return(obj)
 }
+
+#' @importFrom ape branching.times
+#' @importFrom ape is.ultrametric
+add.branch.time.mcmctree <- function(obj, force.ultrametric=FALSE, ...){
+    flag_ultrametric <- is.ultrametric(obj@phylo, option=2) || is.ultrametric(obj@phylo)
+    if (force.ultrametric && flag_ultrametric){
+        message("This tree is not ultrametric, and you has set force.ultrametric to TRUE, so the tree will be converted to ultrametric automatically!")
+        obj <- force.ultrametric(obj)
+    }
+    if (flag_ultrametric){
+        xx <- branching.times(obj@phylo)
+        ntip <- Ntip(obj)
+        N <- Nnode(obj, internal.only = FALSE)
+        xx <- data.frame(node=as.character((ntip+1):N), reltime=as.vector(xx))
+        obj@data <- full_join(obj@data, xx, by="node") 
+    }
+    return(obj)
+}

R/method-as-ultrametric.R

@@ -0,0 +1,48 @@
+#' @title as.ultrametric
+#' @param tree tree object
+#' @param ... additional parameters
+#' @return treedata or phylo object
+#' @export 
+as.ultrametric <- function(tree, ...){
+    UseMethod("as.ultrametric")
+}
+
+#' @method as.ultrametric phylo
+#' @export
+## reference 
+## https://github.com/PuttickMacroevolution/MCMCtreeR/blob/master/R/readMCMCTree.R
+as.ultrametric.phylo <- function(tree, ...){
+    outer <- tree$edge[, 2]
+    inner <- tree$edge[, 1]
+    ntip <- Ntip(tree)
+    totalPath <- c()
+    tipindx <- which(outer <= ntip)
+    for (i in tipindx) {
+        start <- i
+        end <- inner[start]
+        edgeTimes <- tree$edge.length[start]
+        while (end != inner[1]) {
+            start <- which(outer == end)
+            end <- inner[start]
+            edgeTimes <- c(edgeTimes, tree$edge.length[start])
+        }
+        totalPath <- c(totalPath, sum(edgeTimes))
+    }
+    addLength <- max(totalPath) - totalPath
+    tree$edge.length[tipindx] <- tree$edge.length[tipindx] + addLength
+    return (tree)
+}
+
+#' @method as.ultrametric treedata
+#' @export
+as.ultrametric.treedata <- function(tree, ...){
+    tree@phylo <- as.ultrametric(tree=tree@phylo,...)
+    return (tree)
+}
+
+#' @method as.ultrametric tbl_tree
+as.ultrametric.tbl_tree <- function(tree, ...){
+    tree <- as.treedata(tree)
+    tree <- as.ultrametric(tree)
+    return(tree)
+}