From 9337979c7082766e9440ebe79c6e57185cdb42c6 Mon Sep 17 00:00:00 2001
From: mjanuario <januario@umich.edu>
Date: Thu, 24 Aug 2023 15:16:07 -0400
Subject: [PATCH] completely re-wrote draw.sim

---
 R/draw.sim.R | 298 ++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 250 insertions(+), 48 deletions(-)

diff --git a/R/draw.sim.R b/R/draw.sim.R
index 7dc33da..e2334a4 100644
--- a/R/draw.sim.R
+++ b/R/draw.sim.R
@@ -4,10 +4,14 @@
 #' see \code{?sim}) in the graphics window. Allows for the assignment of 
 #' speciation and sampling events, and further customization.
 #'
+#' @inheritParams traits.summary
+#'
 #' @param sim A \code{sim} object, containing extinction times, speciation 
 #' times, parent, and status information for each species in the simulation. 
 #' See \code{?sim}.
 #' 
+#' @traits 
+#' 
 #' @param fossils A \code{data.frame} containing the fossil occurrences of 
 #' each lineage, e.g. as returned by the \code{sample.clade} function. The
 #' format of this argument will define the way fossils are drawn (see below).
@@ -31,7 +35,8 @@
 #' @param lineageColors Character vector giving the colors of all lineages, 
 #' sorted by the original lineage order (the one in the \code{sim} object). 
 #' Must have same length as the number of lineages in the \code{sim} object. 
-#' If NULL (default value) all lineages are plotted as black.
+#' If NULL (default value) all lineages are plotted as black. this parameter 
+#' has no effect if \code{traits} is also provided.
 #' 
 #' @param tipLabels Character vector manually assigning the tip labels of all 
 #' lineages, sorted by the original lineage order (the one in the \code{sim} 
@@ -39,6 +44,36 @@
 #' object. If NULL (default value) all lineages are plotted as "t#", with "#" 
 #' being the position of that lineage in the \code{sim} object.
 #' 
+#' @param traitID Numerical giving the trait which will be plotted. this 
+#' parameter is only useful when multiple traits were simualted in the 
+#' same \code{sim} object.
+#' 
+#' @param traitColors Character vector providing colors providing colors for 
+#' the states of a given trait, so its length must equal or exceed the number 
+#' of states. Default values provide 7 colors (and so they can plot up to 7 
+#' states).
+#' 
+#' @param TraitLegendPlace Placement of state legend. Accepted values are 
+#' \code{"topleft"} (default value), \code{"bottomleft"}, \code{"bottomright"}, 
+#' and \code{"topright"}.
+#' 
+#' @param FossilsToDraw Character assigning if fossils will be represented by 
+#' exact time placements (\code{"exact"}, efault value), by horizontal bars 
+#' giving range information (\code{"ranges"}), or by both forms (\code{"all"}).
+#' 
+#' @param fossilRangeAlpha Numerical giving color transparency for fossil range 
+#' representation. Integers between \code{0} and \code{255} are preferred, 
+#' but any float between \code{0} and \code{1} is also accepted. Default value 
+#' is \code{100}.
+#' 
+#' @param restoreOldPar Logical assigning if plot default values show be 
+#' restored after function finalizes plotting. Deafult is \code{TRUE}, but 
+#' users interesting in using plot additions (e.g. \code{abline()} to highlight 
+#' a certain age) should assign this as \code{FALSE} to use the x and y values 
+#' in the plot. If false, x-axis follows time, and y-axis follows the number of 
+#' species plotted, with 1 being the bottom lineage, and the upper y-limit 
+#' being the Nth lineage in the \code{sim}.
+#' 
 #' @param ... Further arguments to be passed to \code{plot}
 #' 
 #' @return A plot of the simulation in the graphics window. If the 
@@ -54,38 +89,38 @@
 #'
 #' @examples
 #' 
-#' ###
+#' 
 #' # we start drawing a simple simulation
-#'
+#' #'
 #' # maximum simulation time
 #' tMax <- 10
-#'
+#' #'
 #' # set seed
 #' set.seed(1)
-#'
+#' #'
 #' # run a simulation
 #' sim <- bd.sim(n0 = 1, lambda = 0.6, mu = 0.55, tMax = tMax, 
 #'               nFinal = c(10,20)) 
-#'               
+#' 
 #' # draw it
 #' draw.sim(sim)
 #' 
 #' ###
 #' # we can add fossils to the drawing
-#'
+#' #'
 #' # maximum simulation time
 #' tMax <- 10
-#'
+#' #'
 #' # set seed
 #' set.seed(1)
-#'
+#' #'
 #' # run a simulation
 #' sim <- bd.sim(n0 = 1, lambda = 0.6, mu = 0.55, tMax = tMax, 
 #'               nFinal = c(10,20)) 
-#'
+#' #'
 #' # set seed
 #' set.seed(1)
-#'
+#' #'
 #' # simulate data resulting from a fossilization process
 #' # with exact occurrence times
 #' fossils <- sample.clade(sim = sim, rho = 4, tMax = tMax, returnTrue = TRUE)
@@ -120,7 +155,7 @@
 #' 
 #' # set seed
 #' set.seed(20)
-#'
+#' #'
 #' # create time bins randomly
 #' bins <- c(tMax, 0, runif(n = rpois(1, lambda = 6), min = 0, max = tMax))
 #' 
@@ -132,7 +167,8 @@
 #'                         returnTrue = FALSE, bins = bins)
 #' 
 #' # draw it, sorting lineages by their parent
-#' draw.sim(sim, fossils = fossils, sortBy = "PAR")
+#' draw.sim(sim, fossils = fossils, sortBy = "PAR",
+#'          FossilsToDraw = "ranges", restoreOldPar = F)
 #' 
 #' # adding the bounds of the simulated bins
 #' abline(v = bins, lty = 2, col = "blue", lwd = 0.5)
@@ -140,7 +176,7 @@
 #' #alternatively, one can draw lineages varying colors and tip labels:
 #' #(note how they are sorted)
 #' draw.sim(sim, tipLabels = paste0("spp_", 1:10), 
-#'      lineageColors = rep(c("red", "green", "blue"), times=5))
+#'          lineageColors = rep(c("red", "green", "blue"), times=5))
 #' 
 #' ###
 #' # we can control how to sort displayed species exactly
@@ -154,10 +190,10 @@
 #' # run birth-death simulations
 #' sim <- bd.sim(n0 = 1, lambda = 0.6, mu = 0.55, tMax = tMax, 
 #'               nFinal = c(10,20)) 
-#'
+#' #'
 #' # set seed
 #' set.seed(1)  
-#'
+#' #'
 #' # simulate fossil sampling
 #' fossils <- sample.clade(sim = sim, rho = 4, tMax = tMax, returnTrue = TRUE)
 #' 
@@ -165,6 +201,51 @@
 #' # value, for instance)
 #' draw.sim(sim, fossils = fossils, sortBy = sample(1:length(sim$TS)))
 #' 
+#' ###
+#' # plotting trait dependent diversification:
+#' 
+#' # initial number of species
+#' n0 <- 1
+#' 
+#' # maximum simulation time
+#' tMax <- 20
+#' 
+#' # speciation, higher for state 1
+#' lambda <- c(0.1, 0.2)
+#' 
+#' # extinction, lowest for state 0
+#' mu <- c(0.01, 0.03)
+#' 
+#' # number of traits and states (2 binary traits)
+#' nTraits <- 2
+#' nStates <- 2
+#' 
+#' # initial value of both traits
+#' X0 <- 0
+#' 
+#' # transition matrix, with symmetrical transition rates for trait 1,
+#' # and asymmetrical (and higher) for trait 2
+#' Q <- list(matrix(c(0, 0.1,
+#'                    0.1, 0), ncol = 2, nrow = 2),
+#'           matrix(c(0, 1,
+#'                    0.5, 0), ncol = 2, nrow = 2))
+#' 
+#' # set seed
+#' set.seed(1)
+#' 
+#' # run the simulation
+#' sim <- bd.sim.traits(n0, lambda, mu, tMax, nTraits = nTraits, 
+#'                      nStates = nStates, X0 = X0, Q = Q, nFinal = c(2, 10))
+#' 
+#' # maybe we want to take a look at the traits of fossil records too
+#' fossils <- sample.clade(sim$SIM, rho = 0.5, tMax = max(sim$SIM$TS), 
+#'                         returnAll = T,bins = seq(0,20, by=1))
+#' head(fossils)
+#' 
+#' draw.sim(sim$SIM, showLabel = T,  traits = sim$TRAITS, sortBy = "PAR",
+#'          fossils = fossils, FossilsToDraw = "all")
+#' 
+#' 
 #' @importFrom grDevices col2rgb rgb
 #' @importFrom graphics points segments text
 #' 
@@ -172,13 +253,18 @@
 #' @rdname draw.sim
 #' @export
 #' 
-
-draw.sim <- function (sim, fossils = NULL, sortBy = "TS", 
-                      lwdLin = 4, lineageColors=NULL, 
-                      tipLabels=NULL, showLabel = TRUE, ...) {
+draw.sim=function (sim, traits=NULL, fossils = NULL, lineageColors=NULL,
+                   sortBy = "TS", lwdLin = 4, tipLabels=NULL, showLabel = TRUE,
+                   traitID=1, traitColors=c("#a40000", "#16317d", "#007e2f", "#ffcd12", 
+                                            "#b86092", "#721b3e", "#00b7a7"), TraitLegendPlace='topleft',
+                   FossilsToDraw="exact", fossilRangeAlpha=100, restoreOldPar=TRUE, ...) {
+  
   # set up to return par to pre-function settings
   oldPar <- par(no.readonly = TRUE) 
-  on.exit(par(oldPar))
+  
+  if(restoreOldPar){
+    on.exit(par(oldPar))  
+  }
   
   # make NAs 0
   sim$TE[is.na(sim$TE)] <- 0
@@ -211,7 +297,9 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
       min_time <- min(sim$TE, na.rm = TRUE)  
     })
     # if all species are extant, min_time must be adjusted:
-    min_time <- 0
+    if(is.infinite(min_time)){
+      min_time <- 0  
+    }
     xmin <- min_time
     
     if(showLabel){
@@ -240,11 +328,48 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
   
   # fossil data frame
   if (!is.null(fossils)) {
+    
+    if(!(FossilsToDraw %in% c("all", "ranges", "exact"))){
+      stop("\"FossilsToDraw \" must be equal to \"all\", \"ranges\", or to \"exact\"")
+    }
+    
     if (!((c("SampT") %in% colnames(fossils)) | 
           all(c("MaxT", "MinT") %in% colnames(fossils)))) {
       stop("fossils must contain either a SampT or both a MinT and MaxT
            columns. See ?draw.sim and ?sample.clade.")
     }
+    
+    if (!is.null(traits) & !(c("SampT") %in% colnames(fossils))){
+      stop("Plotting fossils with traits require information on true times of fossilization. \n \n Try to create \"fossils\" using \"sample.clade(..., returnAll = T)\"") 
+    }
+  }
+  
+  
+  # trait info:
+  # getting traits:
+  unique_trts <- vector()
+  for(i in 1:length(traits)){
+    unique_trts <- c(unique_trts, unique(traits[[i]][[traitID]]$value))
+  }
+  unique_trts <- unique(unique_trts)
+  unique_trts
+  
+  if(length(traitColors)<length(unique_trts)){
+    stop("Less colors than traits")
+  }
+  
+  # checking length of inputted data:
+  if(!is.null(traits)){
+    unq_lengs <- length(unlist(unique(lapply(list(sim$TE, 
+                                                  sim$TS, sim$PAR, sim$EXTANT, traits), length))))  
+  }else{
+    unq_lengs <- length(unlist(unique(lapply(list(sim$TE, 
+                                                  sim$TS, sim$PAR, sim$EXTANT), length))))
+  }
+  
+  if (unq_lengs != 1) {
+    stop("Trait data does not match with birth-death simulation content. 
+      Please check inputted sim and traits")
   }
   
   # check sortBy
@@ -257,6 +382,17 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
          should have unique indices.")
   }
   
+  #checking fossil input:
+  if(!is.null(fossils)){
+    if(!"SampT" %in% colnames(fossils) & FossilsToDraw != "ranges"){
+      stop("No temporal information to draw requested fossil representation. 
+           \n \n 
+  Try to create \"fossils\" using \"sample.clade(..., returnAll = T)\".
+  \n
+  Alternatively, you might want to set FossilsToDraw = \"ranges\"")
+    }
+  }
+  
   # function for drawing transparency 
   makeTransparent <- function(someColor, alpha = 25) {
     # make color
@@ -303,8 +439,9 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
     ord <- sortBy
   }
   
-  # copy of  sim object in the correct order
+  # copy of sim object in the correct order
   sim_mod <- sim
+  traits_mod <- traits
   
   sim_mod$TE <- sim_mod$TE[ord]
   
@@ -314,26 +451,43 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
   
   sim_mod$EXTANT <- sim_mod$EXTANT[ord]
   
+  traits_mod <-  traits_mod[ord]
+  
   class(sim_mod) <- "sim"
   
   # open plot following user inputs + defaults above
   plot(NA, ...)
   
   # plot durations
-  if(is.null(lineageColors)){
-    lincols <- "black"
+  if(!is.null(traits)){
+    for(i in 1:length(traits_mod)){
+      segments(x0 = traits_mod[[i]][[traitID]]$max, 
+               x1 = traits_mod[[i]][[traitID]]$min, 
+               y1 = i, 
+               y0 = i, 
+               lwd = lwdLin, 
+               # below we need the "+1" bc first trait is "zero"
+               col = traitColors[traits_mod[[i]][[traitID]]$value+1])  
+    }  
   }else{
-    lincols <- lineageColors[ord]
+    if(is.null(lineageColors)){
+      lincols <- "black"
+    }else{
+      lincols <- lineageColors[ord]
+    }
+    
+    segments(x0 = sim_mod$TS, x1 = sim_mod$TE, y1 = 1:length(sim_mod$TE), 
+             y0 = 1:length(sim_mod$TE), lwd = lwdLin, col = lincols) 
   }
   
-  segments(x0 = sim_mod$TS, x1 = sim_mod$TE, y1 = 1:length(sim_mod$TE), 
-           y0 = 1:length(sim_mod$TE), lwd = lwdLin, col = lincols)
+  
+  
   
   # show species labels if user wants
   if (showLabel) {
     if(is.null(tipLabels)){
       tiplabels <- paste0("t", sprintf(paste0("%0", 
-         ceiling(log(length(sim$TE), 10)), "d"), ord))
+                                              ceiling(log(length(sim$TE), 10)), "d"), ord))
     }else{
       tiplabels <- tipLabels[ord]
     }
@@ -341,9 +495,20 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
     labelPosition = sim_mod$TE
     labelPosition[is.na(labelPosition)]=0
     
+    if(!is.null(traits)){
+      textcolsaux = traits.summary(sim, traits = traits)[traitID]  
+      textcols = traitColors[unlist(textcolsaux)+1]
+      tip_traits <- unlist(lapply(traits, function(x) tail(x[[1]]$value, 1)))
+      tip_cols = traitColors[tip_traits+1]
+      tip_cols = tip_cols[ord]
+    }else{
+      tip_cols = rep("black", times=length(sim_mod$TE))
+    }
+    
+    
     text(y = 1:length(sim_mod$TE), 
          x = labelPosition - ((max(sim$TS) - min_time ) * 0.035), 
-         labels = tiplabels)
+         labels = tiplabels, col = tip_cols)
   }
   
   # establish references for plotting
@@ -355,40 +520,77 @@ draw.sim <- function (sim, fossils = NULL, sortBy = "TS",
   aux_y <- unlist(lapply(sim$PAR, function(x) which(ord == x)[1]))
   
   # dashed lines between parent and daughter species
+  if(!is.null(traits)){
+    InheritedTraits = unlist(lapply(traits_mod, function(x) head(x[[traitID]]$value, 1)))
+    vertLinCols = traitColors[InheritedTraits+1]
+  }else{
+    vertLinCols = "gray50"
+  }
+  
+  
   segments(x0 = sim_mod$TS[-luca], x1 = sim_mod$TS[-luca], 
            y1 = (1:length(sim_mod$TE))[-luca], y0 = (aux_y[ord])[-luca], 
-           lty = 2, lwd = lwdLin*0.25, col = "gray50")
+           lty = 2, lwd = lwdLin*0.25, col = vertLinCols)
+  
+  if(!is.null(traits)){
+    legend(TraitLegendPlace, 
+           legend= paste("State ", sort(unique(InheritedTraits)))
+           , col=traitColors, lty=1, lwd=lwdLin)
+    
+  }
   
   # add fossils
   if (!(is.null(fossils))) {
     # find number ID of each species' fossils
     ids <- as.numeric(gsub("t", "", fossils$Species))
     
-    # if SampT is in the data frame, have true occurrence times
-    if ("SampT" %in% colnames(fossils)) {
-      # draw fossil occurrences as time points
-      points(x = fossils$SampT, 
-             y = unlist(lapply(ids, function(x) which(ord == x))), 
-             col = "red", pch = 16, cex = lwdLin*0.25)
+    if(!is.null(traits)){
+      # stablishing traits:
+      traitSummary <- traits.summary(sim, traits, 
+                                     fossils = fossils, selection = "fossil")
       
-      # if MaxT and MinT are also in the columns, message
-      if ("MaxT" %in% colnames(fossils) & "MinT" %in% 
-          colnames(fossils)) {
-        message("fossils contains both SampT and MaxT/MinT columns. Only
-                true fossil occurrences will be drawn.")
-      }
+      colFossils <- traitColors[unlist(traitSummary[traitID])+1]  
+    }else{
+      colFossils <- rep("red", times=length(sim$TE))
     }
-    # if MaxT and MinT are in the data frame, have occurrence time ranges
-    else if ("MaxT" %in% colnames(fossils) & "MinT" %in% 
-             colnames(fossils)) {
+    
+    # if SampT is in the data frame, have true occurrence times
+    if(FossilsToDraw=="ranges"){
+      
       # jitter lines for drawing fossil ranges
       y_jittered <- jitter_foo(unlist(lapply(ids, 
                                              function(x) which(ord == x))))
       
       # draw fossil ranges a bit transparent
       segments(x1 = fossils$MaxT, x0 = fossils$MinT, y1 = y_jittered, 
-               y0 = y_jittered, col = makeTransparent("red", 100), 
-               lwd = lwdLin * 0.75)
+               y0 = y_jittered, col = makeTransparent(someColor = colFossils,
+                                                      alpha = fossilRangeAlpha), lwd = lwdLin * 0.75)
+      
+    }else{
+      if(FossilsToDraw=="exact"){
+        
+        #raw fossil occurrences as time points
+        points(x = fossils$SampT, 
+               y = unlist(lapply(ids, function(x) which(ord == x))), 
+               col = colFossils, pch = 16, cex = lwdLin*0.25)
+        
+      }else{
+        
+        # jitter lines for drawing fossil ranges
+        y_jittered <- jitter_foo(unlist(lapply(ids, 
+                                               function(x) which(ord == x))))
+        
+        # FIRST draw fossil ranges a bit transparent
+        segments(x1 = fossils$MaxT, x0 = fossils$MinT, y1 = y_jittered, 
+                 y0 = y_jittered, col = makeTransparent(someColor = colFossils,
+                                                        alpha = fossilRangeAlpha), lwd = lwdLin * 0.75)
+        
+        # THEN draw fossil occurrences as time points
+        points(x = fossils$SampT, 
+               y = unlist(lapply(ids, function(x) which(ord == x))), 
+               col = colFossils, pch = 16, cex = lwdLin*0.25)
+        
+      }
     }
   }
 }