Merge pull request #169 from PoonLab/dev

Dev

Author: ArtPoon

NAMESPACE

@@ -4,8 +4,11 @@ S3method(as.phylo,InnerTree)
 S3method(as.phylo,OuterTree)
 S3method(plot,Model)
 S3method(plot,OuterTree)
-S3method(plot,twt.counts)
+S3method(plot,dynamics)
+S3method(print,InnerTree)
 S3method(print,Model)
+S3method(print,OuterTree)
+S3method(print,dynamics)
 S3method(summary,Model)
 export(Host)
 export(HostSet)

R/Host.R

@@ -100,6 +100,31 @@ Host <- R6Class(
       } else {
         return(NULL)
       }
+    },
+    
+    print = function() {
+      cat("twt Host", self$get.name(), "\n")
+      cat("  Compartment:", self$get.compartment(), "\n")
+      cat("  Infection time(s):", paste(self$get.transmission.time()), "\n")
+      cat("  Sampling time:", self$get.sampling.time(), "\n")
+      
+      sources <- self$get.source()
+      if (is.list(sources)) {
+        host.names <- sapply(sources, function(h) h$get.name())
+      } else if ("Host" %in% class(sources)) {
+        host.names <- sources$get.name()
+      } else {
+        host.names <- ""
+      }
+      cat("  Source(s):", host.names, "\n")
+      
+      pathogens <- self$get.pathogens()
+      if (is.list(pathogens)) {
+        p.names <- sapply(pathogens, function(p) p$get.name())
+      } else {
+        p.names <- pathogens$get.name()
+      }
+      cat("  Pathogen(s):", p.names, "\n")
     }
   ),
 
@@ -133,6 +158,14 @@ HostSet <- R6Class(
       private$index <- index  # unique IDs for members
     },
 
+    print = function() {
+      my.name <- self$get.name()
+      cat("twt HostSet", ifelse(is.na(my.name), "", my.name), "\n")
+      cat(" ", self$count.type(), "Host objects\n")
+    },
+    
+    get.name = function() { private$name },
+    
     get.names = function() {
       sapply(private$hosts, function(h) { h$get.name() })
     },
@@ -273,3 +306,5 @@ HostSet <- R6Class(
     index = NULL
   )
 )
+
+

R/InnerTree.R

@@ -15,7 +15,7 @@
 InnerTree <- R6Class(
   "InnerTree",
   public = list(
-    initialize = function(outer, mod, prefix='P', p.index=1) {
+    initialize = function(outer, prefix='P', p.index=1) {
       private$inner.log <- data.frame(
         time=numeric(),  # time of event
         event=character(),  # type of event, e.g., coalescence
@@ -27,7 +27,7 @@ InnerTree <- R6Class(
         pathogen2=character()
       )
 
-      private$mod <- mod
+      private$mod <- outer$get.model()  # inherits model from OuterTree
       private$prefix <- prefix
       private$p.index <- p.index
 
@@ -40,7 +40,6 @@ InnerTree <- R6Class(
 
       # inactive Hosts are part of the transmission history
       private$inactive <- HostSet$new()
-      private$inactive$add.host(index.case$clone())
       private$sampled <- HostSet$new()
 
       # note there is overlap between `retired` and `sampled` HostSets
@@ -54,6 +53,13 @@ InnerTree <- R6Class(
         }
       }
 
+      # handle index case
+      if (index.case$get.name() %in% sampled$get.names()) {
+        private$sampled$add.host(index.case$clone())
+      } else {
+        private$inactive$add.host(index.case$clone())
+      }
+      
       # track Hosts with active Pathogen lineages
       private$active <- HostSet$new()
 
@@ -117,7 +123,7 @@ as.phylo.InnerTree <- function(obj) {
   events <- obj$get.log()
   events$time <- as.numeric(events$time)
 
-  active <- inner$get.active()
+  active <- obj$get.active()
   if (active$count.type() != 1) {
     stop("Error, expected only one Host in active HostSet")
   }
@@ -179,7 +185,8 @@ as.phylo.InnerTree <- function(obj) {
     Nnode=length(node.label), edge=edge, 
     edge.length=as.numeric(edge.list$length),
     event=events$event[match(nodes, events$pathogen2)],
-    compartment=edge.list$compartment
+    compartment=edge.list$compartment,
+    host=events$from.host[match(nodes, events$pathogen2)]
   )
   attr(phy, 'class') <- 'phylo'
   phy
@@ -236,7 +243,7 @@ as.phylo.InnerTree <- function(obj) {
 #' @param node:  character, Pathogen name
 #' @param order:  character, 'preorder' or 'postorder'
 #' @param result:  character, vector to append results by recursive calls
-#' @result character, Pathogen names ordered by tree traversal
+#' @return character, Pathogen names ordered by tree traversal
 #' 
 #' @keywords internal
 #' @noRd
@@ -253,3 +260,17 @@ as.phylo.InnerTree <- function(obj) {
   }
   return(result)
 }
+
+
+#' Generic print function for R6 objects of class `InnerTree`
+#' @export
+#' @noRd
+print.InnerTree <- function(obj) {
+  cat("twt InnerTree\n")  # bold color!
+  cat(" ", obj$get.sampled()$count.type(), "sampled Pathogens\n")
+  cat(" ", obj$get.active()$count.type(), "active Pathogens\n")
+  cat(" ", obj$get.inactive()$count.type(), "inactive Pathogens\n")
+  events <- obj$get.log()
+  cat(" ", nrow(events), "events in inner log: ")
+  print(table(events$event))
+} 

R/Model.R

@@ -354,6 +354,29 @@ Model <- R6Class(
           stop("Sampling field must specify `targets` (compartments)")
         }
 
+        # check that targeted compartments have been defined by user
+        if (is.null(private$compartments)) {
+          load.compartments(settings)
+        }
+        for (cn in names(private$sampling$targets)) {
+          if (!is.element(cn, private$compartments)) {
+            stop(paste("In Sampling:targets compartment", cn, 
+                       "has not been declared in Compartments"))
+          }
+          count <- private$sampling$targets[[cn]]
+          if (!is.numeric(count) | count <= 0) {
+            stop(paste("Sampling:targets:", cn, "must be a positive integer"))
+          }
+        }
+        
+      } else if (mode == "fraction") {
+        # contemporaneous sampling occurs at end of simulation time from 
+        # compartments specified by user, to target size or probability
+        # if value is between 0 and 1, then interpret as a probability
+        if (is.null(private$sampling$targets)) {
+          stop("Sampling field must specify `targets` (compartments)")
+        }
+        
         # check that targeted compartments have been defined by user
         if (is.null(private$compartments)) {
           load.compartments(settings)
@@ -379,15 +402,16 @@ Model <- R6Class(
 #' print.Model
 #' S3 class function to display contents of a Model object
 #' @export
+#' @noRd
 print.Model <- function(obj) {
-  cat("twt Model")
-  cat("Parameters:\n")
+  cat("twt Model\n")  # bold
+  cat("  Parameters:\n")
   params <- obj$get.parameters()
   for (key in names(params)) {
     value <- params[[key]]
-    cat("  ", key, ": ", value, "\n")
+    cat("    ", key, ": ", value, "\n")
   }
-  cat("Compartments: ")
+  cat("  Compartments: ")
   compartments <- obj$get.compartments()
   for (cn in compartments) {
     cat(cn, " ")
@@ -399,15 +423,31 @@ print.Model <- function(obj) {
 #' S3 class function to summarize a Model object - simply a wrapper 
 #' around print()
 #' @export
+#' @noRd
 summary.Model <- function(obj) {
   print(obj)
 }
 
 
 #' plot.Model
 #' Plot a graph summarizing compartments and rates
-#' TODO: label edges with rate expressions
+#' @param obj:  R6 object of class `Model`
+#' @param bg:  character, color for uninfected compartments
+#' @param bg2:  character, color for infected compartments
+#' @param mar:  numeric, margin size (default: 1)
+#' @param lwd:  numeric, line width for edges
 #' @export
-plot.Model <- function(obj) {
-  igraph::plot.igraph(obj$get.graph())
+plot.Model <- function(obj, bg='skyblue', bg2='pink2', mar=1, lwd=2) {
+  g <- obj$get.graph()
+  par(mar=rep(mar, 4))
+  igraph::plot.igraph(g, 
+    vertex.size=(nchar(names(V(g)))+2)*8, 
+    vertex.size2=20, 
+    vertex.shape='crectangle', 
+    vertex.color=ifelse(obj$get.infected(), bg2, bg), 
+    vertex.label.family='sans',
+    edge.arrow.size=0.8,
+    edge.arrow.width=1.2,
+    edge.width=lwd
+    )
 }

R/OuterTree.R

@@ -9,6 +9,8 @@ OuterTree <- R6Class(
   "OuterTree",
   public = list(
     initialize = function(mod) {
+      private$model <- mod
+      
       private$outer.log <- data.frame(
         time=numeric(),  # time of event
         event=character(),  # type of event, e.g., migration
@@ -27,6 +29,7 @@ OuterTree <- R6Class(
     },
 
     # accessor functions
+    get.model = function() { private$model },
     get.targets = function() { private$targets }, 
     get.log = function() { private$outer.log },
     get.nrow = function() { nrow(private$outer.log) },
@@ -64,6 +67,7 @@ OuterTree <- R6Class(
   ),
 
   private = list(
+    model = NULL,
     outer.log = NULL,
     targets = NULL,
     sampled = NULL,
@@ -345,3 +349,17 @@ as.phylo.OuterTree <- function(obj) {
   }
   events
 }
+
+
+#' Generic print function for R6 objects of class `OuterTree`
+#' @export
+#' @noRd
+print.OuterTree <- function(obj) {
+  cat("twt OuterTree\n")  # bold color!
+  cat(" ", obj$get.sampled()$count.type(), "sampled Hosts\n")
+  cat(" ", obj$get.active()$count.type(), "active Hosts\n")
+  cat(" ", obj$get.retired()$count.type(), "retired Hosts\n")
+  events <- obj$get.log()
+  cat(" ", nrow(events), "events in outer log:")
+  print(table(events$event))
+}

R/Pathogen.R

@@ -25,6 +25,15 @@ Pathogen <- R6Class(
       private$children <- children
     },
 
+    print = function() {
+      cat("twt Pathogen", self$get.name(), "\n")
+      cat("  Start time:", self$get.start.time(), "\n")
+      cat("  End time:", self$get.end.time(), "\n")
+      cat("  Parent:", self$get.parent()$get.name(), "\n")
+      children <- sapply(self$get.children(), function(child) child$get.name())
+      cat("  Children:", children, "\n")
+    },
+    
     # immutable attributes
     get.name = function() { private$name },
     get.end.time = function() { private$end.time },