R/plot_transtree.R

In dachuwu/DTQbp: A package for modeling case detection, contact tracing and quarantine with branching processes

#' Plot the transmission tree
#' @author Ta-Chou Ng
#'
#' @param df case_data generated by outbreak_model
#' @param max.infector maximum number of infectors shown on the tree
#' @param lyt layout of the transmission tree. can be one of "kk"(defauk=lt),"tree", or ""
#'
#' @return a visNetwork object.
#' @export
#'
#' @importFrom data.table data.table setnames
#' @importFrom igraph graph_from_data_frame layout_as_tree layout_with_kk layout_with_fr
#' @importFrom visNetwork visNetwork visNodes visEdges visIgraphLayout visInteraction visOptions
#'
#' @examples
#'
#'\dontrun{

#'}

plot_transtree <- function(df, max.infector = 50, lyt="kk"){

  df <- df[df$infector%in%c(0:max.infector),]

  df[, `:=`(
    id = as.character(caseid),
    infector_id = as.character(infector))
    ][, titleid :=  ifelse(infector_id=="0",
                          paste0("<b>ID = ", id, "</b><br>",
                                 "Initial seed<br>"),
                          paste0("<b>", id, "</b><br>",
                                 "Infected by ", infector_id, "<br>"))
      ][, title :=  paste0(titleid,
                        "Infected at t = ", round(exposure,3), "<br>",
                        "Asymptomatic = ", ifelse(asym,"<b>Yes</b>", "No"), "<br>",
                        "Onset at t = ", ifelse(asym,"NA",  round(onset,3)), "<br>",
                        "Captured by contact tracing = ", ifelse(missed,"No","<b>Yes</b>"),  "<br>",
                        "Quarantine period = ",  round(quart_time,3),"~",round(quart_end,3), "<br>",
                        "Isolated at t = ",  round(isolated_time,3), "<br>",
                        "Individual R0 = ", R0,  "<br>",
                        "Individual Re = ", Re,  "<br>")
        ][, `:=`(
          color = ifelse(infector_id == "0", "#696969",
                         ifelse(Re > 1, "#DC143C", "#b36200")),
          shape = ifelse(asym, "square", "dot"))]

  reptime <- df$R0-df$Re
  reptime[is.na(reptime)]<-0
  dfp <- data.table::data.table(infector_id = rep(df$id,reptime))[
    infector_id%in%as.character(c(0:max.infector)),]
  dfp[, `:=`(
    id = paste0("p", 1:.N),
    title = paste0("prevented infection by ",infector_id),
    color = "#B0C4DE",
    shape = "dot")]

  ndf <- rbind(df[,c("id","title","color","shape")],
               dfp[,c("id","title","color","shape")])[, label := ""]

  edf <- rbind(df[,c("infector_id","id")],
               dfp[,c("infector_id","id")])[infector_id!="0",]
  data.table::setnames(edf, "infector_id", "from")
  data.table::setnames(edf, "id", "to")
  edf[, color := ifelse(substr(to,1,1)!="p", "#b36200","#B0C4DE")]

  gg <- igraph::graph_from_data_frame(d = edf, vertices = ndf,
                              directed = T)
  if(lyt=="tree"){
    ly <- igraph::layout_as_tree(gg)
    ly[,2] <- -ly[,2]
  } else if(lyt=="kk"){
    ly <- igraph::layout_with_kk(gg)
  } else if(lyt=="fr"){
    ly <- igraph::layout_with_fr(gg)
  } else {
    warning("Unrecognizable lyt. Use default kk layout.")
    ly <- igraph::layout_with_kk(gg)
  }

  G <- visNetwork::visNetwork(nodes = ndf, edges = edf) %>%
        visNetwork::visNodes(size = 12)%>%
        visNetwork::visEdges(arrows = 'to',
                             color = list(opacity = .7))%>%
        visNetwork::visIgraphLayout(layout = "layout.norm", physics = F,
                                    type = "full", layoutMatrix = ly)%>%
        visNetwork::visInteraction(navigationButtons = F,
                                   hover = T,
                                   hideEdgesOnDrag = T,
                                   selectConnectedEdges = T) %>%
        visNetwork::visOptions(highlightNearest = list(enabled=T, algorithm = "hierarchical",
                                                       degree=list(to = 1, from = 10),
                                                       labelOnly=T))
  return(G)
}