Functions_plots_MCMCtreeR.R

# ------------------------------------------------------------------- # 
# Modified function of the MCMCtreeR R function and a helper          # 
# function to help plot the distributions of more than                #
# one clock model together                                            #
#                                                                     #
# Contact: Sandra Alvarez-Carretero <s.alvarez-carretero@qmul.ac.uk>  #
# ------------------------------------------------------------------- # 

# Load library

library( MCMCtreeR )

# Helper function to reorder tips 
# Code taken from StackOverflow to reoder the tips 
# so they are ladderized properly
# https://stackoverflow.com/questions/34364660/how-to-get-correct-order-of-tip-labels-in-ape-after-calling-ladderize-function
reorder_tips <- function( phy ){
  
  # First step is to filter out internal nodes from the the second
  # column of the edge matrix
  phy    <- ladderize(phy, right = TRUE)
  is_tip <- phy$edge[,2] <= length(phy$tip.label)
  ordered_tips <- phy$edge[is_tip, 2]
  
  #Then you can use this vector to extract the tips in the right order:
  phy$tip.label[ordered_tips]
  
  # Return ladderized phy object
  return( phy )
}

# Helper function to plot in the previous tree generated by 
# "mcmc.tree.plot.RETPLOT" the probability densities of the 
# estimated posterior ages for all the nodes using another
# clock model/s 
add.extra.dists <- function (phy = NULL, num.models = 1, last.plot = NULL, 
                             node.ages = NULL, plot.type = "distributions", 
                             all.nodes = NULL, time.correction = 1,
                             density.col = "#00000050", density.border.col = "#00000080",
                             distribution.height = 0.8, transparency = 0.3,
                             return_coords = FALSE ) 
{

  # Define number of extra models 
  if( num.models > 1 ){ 
    
    hpd.ages <- end.age <- node.in.ape <- tip.range <- node.estimates <- 
      vector( mode = "list", length = num.models )
    
    for ( i in seq(1:num.models) ){
      node.ages[[ i ]] <- lapply(node.ages[[ i ]], function(x) x * time.correction)
      phy[[ i ]]$edge.length <- phy[[ i ]]$edge.length * time.correction
      
      hpd.ages[[ i ]]    <- sapply(node.ages[[ i ]], function(xx) HPDinterval(as.mcmc(xx)))
      end.age[[ i ]]     <- hpd.ages[[ i ]][2, 1]
      node.in.ape[[ i ]] <- as.numeric(names(node.ages[[ i ]]))
      
      phy[[ i ]]$root.edge  <- end.age - MCMCtreeR:::nodeTimes(phy[[i]])[1, 1]
      
      if (length(node.ages[[ i ]]) > (Nnode(phy[[ i ]]) + 1)) {
        tip.range[[ i ]] <- hpd.ages[[ i ]][-c(1:(Nnode(phy[[ i ]]) + 1))]
      }
      else {
        tip.range[[ i ]] <- matrix(0, nrow = 2, ncol = Ntip(phy[[ i ]]))
      }
      node.estimates[[ i ]] <- node.ages[[ i ]]
    }
    
  } else if( num.models == 1 ){
    node.ages <- lapply(node.ages, function(x) x * time.correction)
    phy$edge.length <- phy$edge.length * time.correction
    
    hpd.ages    <- sapply(node.ages, function(xx) HPDinterval(as.mcmc(xx)))
    end.age     <- hpd.ages[2, 1]
    node.in.ape <- as.numeric(names(node.ages))
    
    phy$root.edge  <- end.age - MCMCtreeR:::nodeTimes(phy)[1, 1]
    
    if (length(node.ages) > (Nnode(phy) + 1)) {
      tip.range <- hpd.ages[-c(1:(Nnode(phy) + 1))]
    }
    else {
      tip.range <- matrix(0, nrow = 2, ncol = Ntip(phy))
    }
    node.estimates <- node.ages
    
  }

  if (plot.type == "distributions") {
    last.plot.coord <- last.plot
    last.plot.coord <- last.plot
    int <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                   1]]
    ext <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                   2]]
    
    if( num.models > 1 ){ 
      
      len.all.nd <- vector( mode = "list", length = num.models )
      for( i in seq(1:num.models) ){
        all.nodes <- as.numeric(names(node.ages[[ i ]]))
        len.all.nd[[i]] <- length(all.nodes)
        if (length(density.col) < len.all.nd[[i]]) 
          density.col <- rep(density.col, len.all.nd[[i]])
        if (length(density.border.col) < len.all.nd[[i]]) 
          density.border.col <- rep(density.border.col, 
                                    len.all.nd[[i]])
        for (k in 1:( len.all.nd[[i]] ) ) {
          lower.hpd <- hpd.ages[[i]][1, k]
          upper.hpd <- hpd.ages[[i]][2, k]
          node.now  <- node.estimates[[i]][[k]]
          mcmc.hpd  <- node.now[intersect(which(node.now >= 
                                                 lower.hpd), which(node.now <= upper.hpd))]
          mcmc.den <- density(mcmc.hpd)
          xx <- last.plot.coord$xx[1] - mcmc.den$x
          lower.y <- last.plot.coord$yy[all.nodes[[i]][k]]
          upper.y <- max(ext[which(last.plot.coord$edge[, 
                                                        1] == all.nodes[k])])
          scale.den <- (upper.y - lower.y) * distribution.height
          y.scale <- mcmc.den$y * (scale.den/max(mcmc.den$y))
          yy <- y.scale + lower.y
          polygon(xx, yy, col = density.col[k], border = density.border.col[k])
        }
      }
      
    } else if( num.models == 1 ){
      all.nodes <- as.numeric(names(node.ages))
      len.all.nd <- length(all.nodes)
      density.col <- adjustcolor(density.col, alpha.f = transparency) 
      if (length(density.col) < len.all.nd)
        density.col <- rep(density.col, len.all.nd)
      if (length(density.border.col) < len.all.nd) 
        density.border.col <- rep(density.border.col, 
                                  len.all.nd)
      # SAC-240108: Get coords to later add text labels
      coords_k <- vector( mode = "list", len.all.nd )
      for ( k in 1:len.all.nd ) {
        lower.hpd <- hpd.ages[1, k]
        upper.hpd <- hpd.ages[2, k]
        node.now  <- node.estimates[[k]]
        mcmc.hpd  <- node.now[intersect(which(node.now >= 
                                                lower.hpd), which(node.now <= upper.hpd))]
        mcmc.den <- density(mcmc.hpd)
        xx <- last.plot.coord$xx[1] - mcmc.den$x
        lower.y <- last.plot.coord$yy[all.nodes[k]]
        upper.y <- max(ext[which(last.plot.coord$edge[, 
                                                      1] == all.nodes[k])])
        scale.den <- (upper.y - lower.y) * distribution.height
        y.scale <- mcmc.den$y * (scale.den/max(mcmc.den$y))
        yy <- y.scale + lower.y
        polygon(xx, yy, col = density.col[k], border = density.border.col[k])
        # SAC-240108: Add list with x and y values for "k"
        coords_k[[ k ]] <- list( x = xx, y = yy )
      }
      
    }
    
  }
  if( return_coords == TRUE ){
    return( list( coords_k = coords_k ) )
  }
}



# Modified MCMCtreeR::mcmc.tree.plot function. 
# See differences under comments starting with "# SAC ..."
mcmc.tree.plot.RETPLOT <- function (phy = NULL, xlim.scale = NULL, # SAC added argument "xlim.scale"
                                    ladder = TRUE, # SAC added argument "ladder" for ladderize opt
                                    analysis.type = "mcmctree", mcmc.chain = NULL, 
                                    node.ages = NULL, directory.files = NULL, plot.type = "phylogram", 
                                    build.tree = FALSE, node.method = "bar", all.nodes = NULL, 
                                    add.time.scale = TRUE, add.abs.time = TRUE, scale.res = "Epoch", 
                                    label.timescale.names = FALSE, time.correction = 1, col.age = "blue", 
                                    tip.lengths = FALSE, density.col = "#00000050", density.border.col = "#00000080", 
                                    cex.tips = 1, show.tip.label = TRUE, col.tree = "black", 
                                    tip.color = "black", lwd.bar = 1, grey.bars = TRUE, cex.age = 1, 
                                    cex.labels = 1, relative.height = 0.08, n.trees = 2, tip.bar.col = "#ff000050", 
                                    burn.in = 0.25, distribution.height = 0.8, abs.age.mgp = c(3, 0.35, 0),
                                    abs.age.lwd.ticks = 0.7, abs.age.lwd = 0, tck.abs.age = -0.01, 
                                    abs.age.line = -0.4, pos.age = NULL, ...) 
{
  if (is.na(match(tolower(analysis.type), c("mcmctree", "mrbayes", 
                                            "revbayes", "user")))) {
    stop("analysis.type must be one of 'mcmctree', 'mrbayes', 'revbayes', or 'user'")
  }
  if (any(is.na(match(scale.res, c("Eon", "Period", "Epoch", 
                                   "Age"))))) {
    stop("scale.res must be one/combination of 'Eon', 'Period', 'Epoch', or 'Age'")
  }
  mb.tree <- TRUE
  if (is.numeric(match(tolower(analysis.type), "mrbayes"))) 
    mb.tree <- FALSE
  if (mb.tree && plot.type == "distributions" && is.null(mcmc.chain)) 
    stop("The distribution plot.type options needs the full mcmc.chain")
  if (analysis.type == "mcmctree") {
    phy.in <- phy
    phy <- phy.in[[1]]
  }
  if (analysis.type == "mrbayes") {
    phy.data <- read.in.mrbayes(directory.files)
    t.name <- gsub(".con.tre", ".run1.t", directory.files)
    phy <- phy.data$phy
    if (plot.type == "cladogram") 
      stop("sorry plot type 'cladogram' only avilable for mcmctree trees")
  }
  if (analysis.type == "revbayes") {
    phy.data <- read.in.revbayes(directory.files)
    phy <- phy.data$phy
    if (plot.type == "cladogram") 
      stop("sorry plot type 'cladogram' only avilable for mcmctree trees")
  }
  if (analysis.type == "user") {
    if (plot.type == "cladogram") 
      stop("sorry plot type 'cladogram' only avilable for mcmctree trees")
  }
  if (add.time.scale) {
    highest.level <- length(scale.res)
    check.level <- any(is.na(match(scale.res, c("Eon", "Period", 
                                                "Epoch", "Age"))))
    if (is.na(check.level)) 
      stop("scale.res must be one of Eon, Period, Epoch, Age")
    each.height <- -(relative.height * Ntip(phy))/highest.level
    t.depth <- (each.height * highest.level)
    t.height <- Ntip(phy) + (relative.height * Ntip(phy))
    heights <- seq(from = t.depth, to = 0.5, length.out = highest.level + 
                     1)
  }
  else {
    t.depth <- 0
    t.height <- Ntip(phy)
  }
  if (analysis.type == "mcmctree") {
    if (!is.null(mcmc.chain)) {
      internal.node.local <- which(regexpr("t_n", names(mcmc.chain)) != 
                                     -1)
      node.estimates <- mcmc.chain[, internal.node.local] * 
        time.correction
      phy.matrix <- matrix(NA, ncol = 2, nrow = Ntip(phy) * 
                             2 - 2)
      node.in.ape <- as.numeric(gsub("t_n", "", names(node.estimates)))
      hpd.ages <- apply(node.estimates, 2, function(x) HPDinterval(as.mcmc(x)))
      mean.ages <- apply(node.estimates, 2, function(x) mean(x))
      max.ages <- apply(node.estimates, 2, function(x) max(x))
      end.age <- hpd.ages[2, 1]
      if (build.tree) {
        for (ii in 1:length(mean.ages)) phy.matrix[which(node.in.ape[ii] == 
                                                           phy$edge[, 1]), 1] <- mean.ages[ii]
        ext.node <- which(phy$edge[, 2] > Ntip(phy))
        phy.matrix[ext.node, 2] <- mean.ages[-1]
        phy.matrix[-ext.node, 2] <- 0
        b.lengths <- phy.matrix[, 1] - phy.matrix[, 2]
        phy$edge.length <- b.lengths
        phy$root.edge <- end.age - phy.matrix[1, 1]
      }
      else {
        phy$edge.length <- phy$edge.length * time.correction
        phy$root.edge <- end.age - MCMCtreeR:::nodeTimes(phy)[1, 
                                                  1]
      }
      tip.range <- matrix(0, nrow = 2, ncol = Ntip(phy))
    }
    else {
      phy$edge.length <- phy$edge.length * time.correction
      phy.in$nodeAges <- phy.in$nodeAges * time.correction
      hpd.ages <- t(phy.in$nodeAges[, -1])
      mean.ages <- phy.in$nodeAges[, 1]
      end.age <- hpd.ages[2, 1]
      node.in.ape <- as.numeric(rownames(phy.in$nodeAges))
      ext.node <- which(phy$edge[, 2] > Ntip(phy))
      phy$root.edge <- end.age - mean.ages[1]
      tip.range <- matrix(0, nrow = 2, ncol = Ntip(phy))
    }
  }
  if (analysis.type == "mrbayes") {
    node.hpd <- phy.data
    hpd.ages <- sapply(phy.data$node.ages[, "age_95%HPD"], 
                       function(xx) as.numeric(strsplit(xx, ",")[[1]]))
    end.age <- hpd.ages[2, 1]
    node.in.ape <- as.numeric(phy.data$node.ages[, "node.name"])
    phy$root.edge <- end.age - MCMCtreeR:::nodeTimes(phy)[1, 1]
    tip.range <- sapply(phy.data$tip.ages[, "age_95%HPD"], 
                        function(xx) as.numeric(strsplit(xx, ",")[[1]]))
  }
  if (analysis.type == "revbayes") {
    node.hpd <- phy.data
    hpd.ages <- sapply(phy.data$node.ages[, "age_95%_HPD"], 
                       function(xx) as.numeric(strsplit(xx, ",")[[1]]))
    end.age <- hpd.ages[2, 1]
    node.in.ape <- as.numeric(phy.data$node.ages[, "node.name"])
    phy$root.edge <- end.age - MCMCtreeR:::nodeTimes(phy)[1, 1]
    if (colnames(phy.data$tip.ages) == "age_95%_HPD") {
      tip.range <- sapply(phy.data$tip.ages[, "age_95%_HPD"], 
                          function(xx) as.numeric(strsplit(xx, ",")[[1]]))
    }
    else {
      tip.range <- matrix(0, nrow = 2, ncol = Ntip(phy))
    }
  }
  if (analysis.type == "user") {
    node.ages <- lapply(node.ages, function(x) x * time.correction)
    phy$edge.length <- phy$edge.length * time.correction
    hpd.ages <- sapply(node.ages, function(xx) HPDinterval(as.mcmc(xx)))
    end.age <- hpd.ages[2, 1]
    node.in.ape <- as.numeric(names(node.ages))
    phy$root.edge <- end.age - MCMCtreeR:::nodeTimes(phy)[1, 1]
    if (length(node.ages) > (Nnode(phy) + 1)) {
      tip.range <- hpd.ages[-c(1:(Nnode(phy) + 1))]
    }
    else {
      tip.range <- matrix(0, nrow = 2, ncol = Ntip(phy))
    }
    node.estimates <- node.ages
  }
  if (plot.type == "phylogram") {
    plot.phylo(ladderize(phy, right = ladder), type = "phylogram", edge.col = col.tree, 
               root.edge = TRUE, y.lim = c(t.depth, t.height), cex = cex.tips, 
               show.tip.label = show.tip.label, tip.color = tip.color, 
               ...)
    last.plot.coord <- get("last_plot.phylo", envir = .PlotPhyloEnv)
    int <- last.plot.coord$yy[last.plot.coord$edge[, 1]]
    ext <- last.plot.coord$yy[last.plot.coord$edge[, 2]]
    structure <- cbind(int, ext)
    node.in.ape <- c(node.in.ape, 1:Ntip(phy))
    mean.ages <- cbind(end.age - hpd.ages, end.age - tip.range)
    match.nodes <- last.plot.coord$edge[1, 1]
    location.here <- match(match.nodes, node.in.ape)
    xx <- mean.ages[, location.here]
    match.edges <- which(last.plot.coord$edge[, 1] == Ntip(phy) + 
                           1)
    if (node.method != "none") {
      if (is.null(all.nodes)) {
        ext.node <- which(last.plot.coord$edge[, 2] > 
                            Ntip(phy))
      }
      else {
        ext.node <- match(all.nodes, last.plot.coord$edge[, 
                                                          2])
      }
      if (length(col.age) < length(ext.node)) 
        col.age <- rep(col.age, length(ext.node) + 1)
      if (node.method == "full.length") {
        polygon(c(xx[1], xx[2], xx[2], xx[1]), matrix(c(0, 
                                                        0, Ntip(phy), Ntip(phy)), ncol = 2), col = col.age[1], 
                border = FALSE)
      }
      if (node.method == "node.length") {
        polygon(c(xx[1], xx[2], xx[2], xx[1]), rbind(structure[match.edges, 
                                                               2], structure[match.edges, 2]), col = col.age[1], 
                border = FALSE)
      }
      if (node.method == "bar") {
        lines(c(xx[1], xx[2]), rep(structure[1, 1], 2), 
              col = col.age[1], lwd = lwd.bar)
      }
      if (length(col.age) == length(ext.node)) 
        col.age <- rep(col.age, length(ext.node) + 1)
      counter <- 2
      for (uu in ext.node) {
        match.nodes <- last.plot.coord$edge[uu, 2]
        location.here <- match(match.nodes, node.in.ape)
        xx <- mean.ages[, location.here]
        match.edges <- range(which(last.plot.coord$edge[, 
                                                        1] == last.plot.coord$edge[uu, 2]))
        if (node.method == "full.length") {
          polygon(c(xx[1], xx[2], xx[2], xx[1]), matrix(c(0, 
                                                          0, Ntip(phy), Ntip(phy)), ncol = 2), col = col.age[counter], 
                  border = FALSE)
        }
        if (node.method == "node.length") {
          polygon(c(xx[1], xx[2], xx[2], xx[1]), rbind(structure[match.edges, 
                                                                 2], structure[match.edges, 2]), col = col.age[counter], 
                  border = FALSE)
        }
        if (node.method == "bar") {
          lines(xx, rep(structure[uu, 2], 2), col = col.age[counter], 
                lwd = lwd.bar)
        }
        counter = counter + 1
      }
    }
    if (tip.lengths) {
      phy.to <- MCMCtreeR:::nodeTimes(phy)[, 2]
      tip.here <- which(phy$edge[, 2] <= Ntip(phy))
      tip.to.label <- tip.here[which(MCMCtreeR:::nodeTimes(phy)[tip.here, 
                                                    2] > 0.01)]
      for (uu in phy$edge[tip.to.label, 2]) {
        match.nodes <- Nnode(phy) + uu
        xx <- mean.ages[, match.nodes]
        match.edges <- last.plot.coord$edge[uu, 1]
        yy.loc <- last.plot.coord$yy[uu]
        lines(xx, rep(yy.loc, 2), col = tip.bar.col, 
              lwd = lwd.bar)
      }
    }
  }
  if (plot.type == "cladogram") {
    plot.phylo(ladderize(phy, right = ladder), type = "cladogram", edge.col = col.tree, 
               root.edge = TRUE, y.lim = c(t.depth, t.height), cex = cex.tips, 
               show.tip.label = show.tip.label, tip.color = tip.color, 
               ...)
    last.plot.coord <- get("last_plot.phylo", envir = .PlotPhyloEnv)
    int <- last.plot.coord$yy[last.plot.coord$edge[, 1]]
    ext <- last.plot.coord$yy[last.plot.coord$edge[, 2]]
    structure <- cbind(int, ext)
    node.in.ape <- c(node.in.ape, 1:Ntip(phy))
    mean.ages <- cbind(end.age - hpd.ages, end.age - tip.range)
    match.nodes <- last.plot.coord$edge[1, 1]
    location.here <- match(match.nodes, node.in.ape)
    xx <- mean.ages[, location.here]
    match.edges <- which(last.plot.coord$edge[, 1] == Ntip(phy) + 
                           1)
    match.nodes <- last.plot.coord$edge[1, 1]
    location.here <- which(match.nodes == last.plot.coord$edge[, 
                                                               1])
    xx <- mean.ages[, 1]
    ext.node <- which(last.plot.coord$edge[, 2] > Ntip(phy))
    if (length(col.age) < dim(last.plot.coord$edge)[1]) {
      col.age <- rep(col.age, dim(last.plot.coord$edge)[1])
      lines(xx, rep(structure[1, 1], 2), lwd = lwd.bar, 
            col = col.age[1])
    }
    counter <- 1
    for (uu in 1:dim(last.plot.coord$edge)[1]) {
      match.nodes <- last.plot.coord$edge[uu, ]
      location.here <- match(match.nodes, node.in.ape)
      xx <- c(mean.ages[, location.here])
      polygon(c(xx[1], xx[3], xx[4], xx[2]), c(structure[uu, 
                                                         ], rev(structure[uu, ])), col = col.age[counter], 
              border = FALSE)
      counter <- counter + 1
    }
  }
  if (plot.type == "distributions") {
    # SAC added the option of passing xlim.scale in case 
    # the user wants to plot a specific range of x.axis to 
    # capture future extra distributions from other models 
    # that lie over the ones limited by current model
    # Now and if,else conditional takes care of this
    # I also added the option "ladder" depending on the
    # user's decision to have it ladderized or not - default = TRUE
    if( is.null(xlim.scale) == FALSE ){
      plot.phylo(ladderize(phy, right = ladder), type = "p", edge.col = col.tree, 
               root.edge = TRUE, y.lim = c(t.depth, t.height), cex = cex.tips, 
               show.tip.label = show.tip.label, tip.color = tip.color,
               x.lim = xlim.scale, 
               ...)
    }else{
      plot.phylo(ladderize(phy, right = ladder), type = "p", edge.col = col.tree, 
                 root.edge = TRUE, y.lim = c(t.depth, t.height), cex = cex.tips, 
                 show.tip.label = show.tip.label, tip.color = tip.color,
                 ...)
    }
    last.plot.coord <- get("last_plot.phylo", envir = .PlotPhyloEnv)
    last.plot.coord <- get("last_plot.phylo", envir = .PlotPhyloEnv)
    if (analysis.type == "mrbayes") {
      internal.nodes <- unique(phy$edge[, 1])
      node.times <- vector("list", length(internal.nodes))
      for (kk in 1:n.trees) {
        tre.in.ape <- read.nexus(gsub("run1", paste0("run", 
                                                     kk), t.name))
        tre.in <- scan(gsub("run1", paste0("run", kk), 
                            t.name), what = "", sep = "\n", quiet = TRUE)
        start.names <- which(regexpr("translate", tre.in) != 
                               -1) + 1
        end.names <- which(regexpr("tree gen", tre.in) != 
                             -1)[1] - 1
        names.on.tree <- sapply(tre.in[start.names:end.names], 
                                function(xx) {
                                  xx <- gsub(",", "", tail(strsplit(xx, " ")[[1]], 
                                                           1))
                                  xx <- gsub(";", "", xx)
                                })
        clock.rate.in <- unlist(gregexpr("&clockrate", 
                                         tre.in))
        clock.rate.loc <- which(clock.rate.in != -1)
        ceiling.rm <- ceiling(length(tre.in.ape) * burn.in)
        tre.in.ape <- tre.in.ape[-c(1:ceiling.rm)]
        clock.rate.loc <- clock.rate.loc[-c(1:ceiling.rm)]
        for (y in 1:length(tre.in.ape)) {
          tree.here <- tre.in[clock.rate.loc[y]]
          new.string <- substring(tree.here, regexpr("&clockrate", 
                                                     tree.here) + 1 )
          new.string <- substring(new.string, 1, regexpr("[]]", 
                                                         new.string) - 1)
          clock.rate <- as.numeric(strsplit(new.string, 
                                            "=")[[1]][2])
          tree.in <- read.tree(text = tree.here)
          tree.in$tip.label <- as.character(names.on.tree[as.numeric(tree.in$tip.label)])
          tree.edge.unique <- unique(tree.in$edge[, 1])
          node.times.now <- MCMCtreeR:::nodeTimes(tree.in)/clock.rate
          parts <- prop.part(tree.in, phy)
          match.on.tree <- which(attributes(parts)$number == 
                                   2)
          for (ii in 1:length(match.on.tree)) {
            name.now <- parts[[match.on.tree[ii]]]
            five <- match(tree.edge.unique[match.on.tree[ii]], 
                          tree.in$edge[, 1])
            node.time.add <- node.times.now[five, 1]
            names.now <- attributes(parts)$labels[name.now]
            mrca.now <- getMRCA(phy, names.now)
            mrca.store <- match(mrca.now, internal.nodes)
            node.times[[mrca.store]] <- c(node.times[[mrca.store]], 
                                          node.time.add)
          }
        }
      }
      names(node.times) <- internal.nodes
      all.nodes <- (Ntip(phy) + 1):(Nnode(phy) + Ntip(phy))
      int <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     1]]
      ext <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     2]]
      len.all.nd <- length(all.nodes)
      if (length(density.col) < len.all.nd) 
        density.col <- rep(density.col, len.all.nd)
      if (length(density.border.col) < len.all.nd) 
        density.border.col <- rep(density.border.col, 
                                  len.all.nd)
      for (k in 1:len.all.nd) {
        location.here <- match(as.numeric(names(node.times)[k]), 
                               node.in.ape)
        lower.hpd <- hpd.ages[1, location.here]
        upper.hpd <- hpd.ages[2, location.here]
        mcmc.hpd <- node.times[[k]][intersect(which(node.times[[k]] >= 
                                                      lower.hpd), which(node.times[[k]] <= upper.hpd))]
        mcmc.den <- density(mcmc.hpd)
        xx <- last.plot.coord$xx[1] - mcmc.den$x
        lower.y <- last.plot.coord$yy[all.nodes[k]]
        upper.y <- max(ext[which(all.nodes[k] == last.plot.coord$edge[, 
                                                                      1])])
        scale.den <- (upper.y - lower.y) * distribution.height
        y.scale <- mcmc.den$y * (scale.den/max(mcmc.den$y))
        yy <- y.scale + lower.y
        polygon(xx, yy, col = density.col[k], border = density.border.col[k])
      }
    }
    if (analysis.type == "revbayes") {
      internal.nodes <- unique(phy$edge[, 1])
      node.times <- vector("list", length(internal.nodes))
      slash <- max(gregexpr("/", directory.files)[[1]])
      path.to.files <- substring(directory.files, 1, slash)
      log.files <- list.files(path = path.to.files, pattern = ".trees")
      n.trees <- length(log.files)
      for (kk in 1:n.trees) {
        tre.in.ape <- read.tree(paste0(path.to.files, 
                                       "/", log.files[kk]))
        ceiling.rm <- ceiling(length(tre.in.ape) * burn.in)
        tre.in.ape <- tre.in.ape[-c(1:ceiling.rm)]
        for (y in 1:length(tre.in.ape)) {
          tree.in <- tre.in.ape[[y]]
          tree.edge.unique <- unique(tree.in$edge[, 1])
          node.times.now <- MCMCtreeR:::nodeTimes(tree.in)
          parts <- prop.part(tree.in, phy)
          match.on.tree <- which(attributes(parts)$number == 
                                   2)
          for (ii in 1:length(match.on.tree)) {
            name.now <- parts[[match.on.tree[ii]]]
            five <- match(tree.edge.unique[match.on.tree[ii]], 
                          tree.in$edge[, 1])
            node.time.add <- node.times.now[five, 1]
            names.now <- attributes(parts)$labels[name.now]
            mrca.now <- getMRCA(phy, names.now)
            mrca.store <- match(mrca.now, internal.nodes)
            node.times[[mrca.store]] <- c(node.times[[mrca.store]], 
                                          node.time.add)
          }
        }
      }
      names(node.times) <- internal.nodes
      all.nodes <- (Ntip(phy) + 1):(Nnode(phy) + Ntip(phy))
      int <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     1]]
      ext <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     2]]
      len.all.nd <- length(all.nodes)
      if (length(density.col) < len.all.nd) 
        density.col <- rep(density.col, len.all.nd)
      if (length(density.border.col) < len.all.nd) 
        density.border.col <- rep(density.border.col, 
                                  len.all.nd)
      for (k in 1:len.all.nd) {
        location.here <- match(as.numeric(names(node.times)[k]), 
                               node.in.ape)
        lower.hpd <- hpd.ages[1, location.here]
        upper.hpd <- hpd.ages[2, location.here]
        mcmc.hpd <- node.times[[k]][intersect(which(node.times[[k]] >= 
                                                      lower.hpd), which(node.times[[k]] <= upper.hpd))]
        mcmc.den <- density(mcmc.hpd)
        xx <- last.plot.coord$xx[1] - mcmc.den$x
        lower.y <- last.plot.coord$yy[all.nodes[k]]
        upper.y <- max(ext[which(all.nodes[k] == last.plot.coord$edge[, 
                                                                      1])])
        scale.den <- (upper.y - lower.y) * distribution.height
        y.scale <- mcmc.den$y * (scale.den/max(mcmc.den$y))
        yy <- y.scale + lower.y
        polygon(xx, yy, col = density.col[k], border = density.border.col[k])
      }
    }
    if (analysis.type == "mcmctree") {
      all.nodes <- (Ntip(phy) + 1):(Nnode(phy) + Ntip(phy))
      int <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     1]]
      ext <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     2]]
      len.all.nd <- length(all.nodes)
      if (length(density.col) < len.all.nd) 
        density.col <- rep(density.col, len.all.nd)
      if (length(density.border.col) < len.all.nd) 
        density.border.col <- rep(density.border.col, 
                                  len.all.nd)
      for (k in 1:len.all.nd) {
        lower.hpd <- hpd.ages[1, k]
        upper.hpd <- hpd.ages[2, k]
        mcmc.hpd <- node.estimates[, k][intersect(which(node.estimates[, 
                                                                       k] >= lower.hpd), which(node.estimates[, k] <= 
                                                                                                 upper.hpd))]
        mcmc.den <- density(mcmc.hpd)
        xx <- last.plot.coord$xx[1] - mcmc.den$x
        lower.y <- last.plot.coord$yy[all.nodes[k]]
        upper.y <- max(ext[which(all.nodes[k] == last.plot.coord$edge[, 
                                                                      1])])
        scale.den <- (upper.y - lower.y) * distribution.height
        y.scale <- mcmc.den$y * (scale.den/max(mcmc.den$y))
        yy <- y.scale + lower.y
        polygon(xx, yy, col = density.col[k], border = density.border.col[k])
      }
    }
    if (analysis.type == "user") {
      all.nodes <- as.numeric(names(node.ages))
      int <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     1]]
      ext <- last.plot.coord$yy[last.plot.coord$edge[, 
                                                     2]]
      len.all.nd <- length(all.nodes)
      if (length(density.col) < len.all.nd) 
        density.col <- rep(density.col, len.all.nd)
      if (length(density.border.col) < len.all.nd) 
        density.border.col <- rep(density.border.col, 
                                  len.all.nd)
      for (k in 1:len.all.nd) {
        lower.hpd <- hpd.ages[1, k]
        upper.hpd <- hpd.ages[2, k]
        node.now <- node.estimates[[k]]
        mcmc.hpd <- node.now[intersect(which(node.now >= 
                                               lower.hpd), which(node.now <= upper.hpd))]
        mcmc.den <- density(mcmc.hpd)
        xx <- last.plot.coord$xx[1] - mcmc.den$x
        lower.y <- last.plot.coord$yy[all.nodes[k]]
        upper.y <- max(ext[which(last.plot.coord$edge[, 
                                                      1] == all.nodes[k])])
        scale.den <- (upper.y - lower.y) * distribution.height
        y.scale <- mcmc.den$y * (scale.den/max(mcmc.den$y))
        yy <- y.scale + lower.y
        polygon(xx, yy, col = density.col[k], border = density.border.col[k])
      }
    }
  }
  if (add.time.scale) {
    
    xlimit <- last.plot.coord$xx
    total.length <- max.xlim <- max(xlimit)
    # SAC changed the default way of looking for Eon, Period, Epoch, and
    # Age variables. Now the code uses MCMCtreeR:::Eon, MCMCtreeR:::Period,
    # MCMCtreeR:::Epoch, and MCMCtreeR:::Age to find them.
    eon.bin <- .bincode(max.xlim, MCMCtreeR:::Eon)
    #>> 240119-SAC Update period ages for Hadean and Eoarchean
    #>> All modifications below when `MCMCtreeR:::Period` appears have been
    #>> changed!
    # I need to create an additional "bin" for Hadean because, otherwise,
    # I cannot reach the limit (max.xlim is 4337.449, and so I need to 
    # extend Hadean until 4600 to allow for this bin to appear!)
    # Each bin signals the "starting point", not the end, that's why.
    # The last bin, pos 21, is the end of Hadean and pos 20 is the start of
    # Hadean
    period_mcmctreeR <- rep( 0, 21 )
    period_mcmctreeR[c(1:20)] <- MCMCtreeR:::Period
    names( period_mcmctreeR )[c(1:20)] <- names( MCMCtreeR:::Period )
    period_mcmctreeR[c(19,20)] <- c( 3600, 4031 )
    period_mcmctreeR[c(21)] <- 4600
    names( period_mcmctreeR )[21] <- "Hadean"
    # cat( "max.xlim:", max.xlim, "\ntweaked:", period_mcmctreeR,
    #      "original:",  MCMCtreeR:::Period, "\n" )
    #period.bin <- .bincode(max.xlim, MCMCtreeR:::Period)
    period.bin <- .bincode(max.xlim, period_mcmctreeR)
    epoch.bin <- .bincode(max.xlim, MCMCtreeR:::Epoch)
    if (is.na(epoch.bin)) 
      epoch.bin <- length(MCMCtreeR:::Epoch)
    age.bin <- .bincode(max.xlim, MCMCtreeR:::Age)
    if (is.na(age.bin)) 
      age.bin <- length(MCMCtreeR:::Age)
    bin.eon <- max.xlim - c(MCMCtreeR:::Eon[1:eon.bin], max.xlim)
    #bin.period <- max.xlim - c(MCMCtreeR:::Period[1:period.bin], max.xlim)
    bin.period <- max.xlim - c(period_mcmctreeR[1:period.bin], max.xlim)
    bin.epoch <- max.xlim - c(MCMCtreeR:::Epoch[1:epoch.bin], max.xlim)
    bin.age <- max.xlim - c(MCMCtreeR:::Age[1:age.bin], max.xlim)
    plot.bar <- function(bin.data, data.colour, name.data, 
                         ht, ht.pl, grey.line = FALSE, cex.labels.int) {
      names.in <- name.data[1:(length(bin.data) - 1)]
      for (u in 2:length(bin.data)) {
        polygon(rep(bin.data[(u):(u - 1)], each = 2), 
                c(heights[ht:ht.pl], rev(heights[ht:ht.pl])), 
                xpd = TRUE, col = data.colour[c(u - 1)], border = FALSE)
        bin.length <- abs(diff(bin.data[(u):(u - 1)]))
        percent <- bin.length/total.length
        if (percent >= 0.4) {
          text(mean(bin.data[(u):(u - 1)]), mean(heights[ht:ht.pl]), 
               names.in[c(u - 1)], cex = cex.labels.int)
        }
        if (percent >= 0.1 && percent < 0.4) {
          name.in <- paste0(paste0(strsplit(names.in[c(u - 
                                                         1)], "")[[1]][1:2], collapse = ""), ".")
          text(mean(bin.data[(u):(u - 1)]), mean(heights[ht:ht.pl]), 
               name.in, cex = cex.labels.int)
        }
        if (percent < 0.1 && percent >= 0.04) {
          name.in <- paste0(paste0(strsplit(names.in[c(u - 
                                                         1)], "")[[1]][1], collapse = ""), ".")
          text(mean(bin.data[(u):(u - 1)]), mean(heights[ht:ht.pl]), 
               name.in, cex = cex.labels.int)
        }
        if (grey.line) {
          if (u%%2 == 0) 
            polygon(rep(bin.data[(u):(u - 1)], each = 2), 
                    c(0.5, Ntip(phy), Ntip(phy), 0.5), xpd = TRUE, 
                    col = "#00000020", border = FALSE)
        }
      }
    }
    plot.time.scale <- function(level, start.one, end.one, 
                                time.check = FALSE, grey.now = FALSE, cex.age.int) {
      if (level == "Eon") {
        plot.bar(bin.eon, MCMCtreeR:::Eon.colour, names(MCMCtreeR:::Eon.colour), 
                 start.one, end.one, grey.line = grey.now, cex.labels.int = cex.labels)
        if (label.timescale.names) 
          text(total.length, mean(heights[start.one:end.one]), 
               "Eon", pos = 4, font = 4, cex = cex.labels)
        if (time.check) {
          if (is.null(pos.age)) 
            pos.age <- Ntip(phy) + (Ntip(phy)/100)
          axis(3, at = bin.eon, labels = signif(total.length - 
                                                  bin.eon, 2), line = abs.age.line, lwd = abs.age.lwd, 
               lwd.ticks = abs.age.lwd.ticks, mgp = abs.age.mgp, 
               cex.axis = cex.age.int, pos = pos.age, tck = tck.abs.age)
          if (label.timescale.names) 
            text(total.length, Ntip(phy) + (Ntip(phy)/35), 
                 "Ma", pos = 4, font = 4, cex = cex.labels, 
                 adj = c(0, 1), xpd = TRUE)
        }
      }
      if (level == "Period") {
        plot.bar(bin.period, MCMCtreeR:::Period.colour, names(MCMCtreeR:::Period.colour), 
                 start.one, end.one, grey.line = grey.now, cex.labels.int = cex.labels)
        if (label.timescale.names) 
          text(total.length, mean(heights[start.one:end.one]), 
               "Period", pos = 4, font = 4, cex = cex.labels)
        if (time.check) {
          if (is.null(pos.age)) 
            pos.age <- Ntip(phy) + (Ntip(phy)/100)
          axis(3, at = bin.period, labels = signif(total.length - 
                                                     bin.period, 2), line = abs.age.line, lwd = abs.age.lwd, 
               lwd.ticks = abs.age.lwd.ticks, mgp = abs.age.mgp, 
               cex.axis = cex.age.int, pos = pos.age, tck = tck.abs.age)
          if (label.timescale.names) 
            text(total.length, Ntip(phy) + (Ntip(phy)/35), 
                 "Ma", pos = 4, font = 4, cex = cex.labels, 
                 adj = c(0, 1), xpd = TRUE)
        }
      }
      if (level == "Epoch") {
        plot.bar(bin.epoch, MCMCtreeR:::Epoch.colour, names(MCMCtreeR:::Epoch.colour), 
                 start.one, end.one, grey.line = grey.now, cex.labels.int = cex.labels)
        if (label.timescale.names) 
          text(total.length, mean(heights[start.one:end.one]), 
               "Epoch", pos = 4, font = 4, cex = cex.labels)
        if (time.check) {
          if (is.null(pos.age)) 
            pos.age <- Ntip(phy) + (Ntip(phy)/100)
          axis(3, at = bin.epoch, labels = signif(total.length - 
                                                    bin.epoch, 2), line = abs.age.line, lwd = abs.age.lwd, 
               lwd.ticks = abs.age.lwd.ticks, mgp = abs.age.mgp, 
               cex.axis = cex.age.int, pos = pos.age, tck = tck.abs.age, 
               ...)
          if (label.timescale.names) 
            text(total.length, Ntip(phy) + (Ntip(phy)/35), 
                 "Ma", pos = 4, font = 4, cex = cex.labels, 
                 adj = c(0, 1), xpd = TRUE)
        }
      }
      if (level == "Age") {
        plot.bar(bin.age, MCMCtreeR:::Age.colour, names(MCMCtreeR:::Age.colour), 
                 start.one, end.one, grey.line = grey.now, cex.labels.int = cex.labels)
        if (label.timescale.names) 
          text(total.length, mean(heights[start.one:end.one]), 
               "Age", pos = 4, font = 4, cex = cex.labels)
        if (time.check) {
          if (is.null(pos.age)) 
            pos.age <- Ntip(phy) + (Ntip(phy)/100)
          axis(3, at = bin.age, labels = signif(total.length - 
                                                  bin.age, 2), line = abs.age.line, lwd = abs.age.lwd, 
               lwd.ticks = abs.age.lwd.ticks, mgp = abs.age.mgp, 
               cex.axis = cex.age.int, pos = pos.age, tck = tck.abs.age)
          if (label.timescale.names) 
            text(total.length, Ntip(phy) + (Ntip(phy)/35), 
                 "Ma", pos = 4, font = 4, cex = cex.labels, 
                 adj = c(0, 1), xpd = TRUE)
        }
      }
    }
    start.one <- 1
    end.one <- 2
    if (length(add.abs.time) == 1 && add.abs.time[1] == TRUE) {
      add.abs.time <- rep(FALSE, length(scale.res))
      add.abs.time[length(add.abs.time)] <- TRUE
    }
    if (length(add.abs.time) == 1 && add.abs.time[1] == FALSE) {
      add.abs.time <- rep(FALSE, length(scale.res))
    }
    if (grey.bars) {
      grey.now.int <- rep(FALSE, length(scale.res))
      if (!any(add.abs.time)) {
        add.time <- length(scale.res)
      }
      else {
        add.time <- which(add.abs.time)
      }
      grey.now.int[add.time] <- TRUE
      for (uu in 1:length(scale.res)) {
        plot.time.scale(scale.res[uu], time.check = add.abs.time[uu], 
                        start.one = start.one, end.one = end.one, grey.now = grey.now.int[uu], 
                        cex.age.int = cex.age)
        print( grey.now.int[uu] )
      
        start.one <- start.one + 1
        end.one <- end.one + 1
      }
    # SAC added this "else" to the code because
    # otherwise it cannot consider when the user
    # does not want to have grey bars and they 
    # keep being plotted despite saying grey.bars = FALSE.
    # In this else conditional, grey.now = FALSE
    # and the user can really get a plot without bars
    }else {
      for (uu in 1:length(scale.res)) {
        plot.time.scale(scale.res[uu], time.check = add.abs.time[uu], 
                        start.one = start.one, end.one = end.one, grey.now = FALSE, 
                        cex.age.int = cex.age)
        start.one <- start.one + 1
        end.one <- end.one + 1
      }
      
    }
  }
  # SAC added this return function so 
  # the resulting coordinates for the plot can be later used
  # by the helper function "add.extra.dists". Therefore,
  # the output of this function is not only a plot but also 
  # a list of objects that can be saved in an object to later
  # be used by add.extra.dists.
  return( last.plot.coord )
}