R/viz_chains.R

In mrajeev08/treerabid: Reconstruct Rabies Transmission Trees

#' Turn a transmssion tree into a graph object
#'
#' @param ttree data.table with columns `id_progen` (id of case progenitor),
#' `id_case` (id of case), `t` (date of case), `x_coord`, and `y_coord`
#' @param from progenitor id's
#' @param to case id's
#' @param attrs data.table of vertex attributes where first column is the
#'  vertex name (i.e. corresponding to unique case ids); if passing to
#'  chain stats, this data.table should atleast have one additional column
#'  named `t` which has the date of the case. Cannot have duplicate vertex names!
#' @return a graph generated by igraph
#' @export
#'
get_graph <- function(from, to, attrs) {

  I_gr <- data.table(from = from[!is.na(from)], to = to[!is.na(from)])
  gr <- igraph::graph_from_data_frame(d = I_gr,
                              vertices = attrs,
                              directed = TRUE)
  igraph::V(gr)$membership <- igraph::components(gr)$membership
  return(gr)
}


#' Get chain size & lengths
#'
#' @param gr graph of transmission tree
#'
#' @return a data.table with membership (i.e. chain id), chain size, and chain length
#' @export
#'
get_chain_stats <- function(gr) {

  comps <- igraph::components(gr)
  return(data.table(membership = seq(comps$no), size = comps$csize,
                    length = unlist(lapply(igraph::decompose(gr), igraph::diameter))))
}


#' Get membership (i.e. chain id) for each case
#'
#' @inheritParams get_chain_stats
#'
#' @return a data.table with id_case, membership, and t_days (time in days from
#'  the start date, i.e. the earliest case in the transmission tree)
#' @export
#'
get_chain_membership <- function(gr) {

  return(
    data.table(membership = igraph::vertex_attr(gr, "membership"),
               id_case = igraph::vertex_attr(gr, "name"),
               t_days = igraph::vertex_attr(gr, "t"))
  )
}


#' Get duration of chain persistence in days
#'
#' @param chain_membership output from `get_chain_membership` (a data.table
#'  with the membership of each case and the time in days)
#'
#' @return a data.table with the start and end days (from the origin, i.e. minimum
#'  case date in the transmission tree) and the duration of the chain in days.
#' @export
#'
get_chain_persistence <- function(chain_membership) {

  chain_pers <- chain_membership[, .(start_date = min(t_days), end_date = max(t_days)),
                                 by = "membership"]
  chain_pers[, days := end_date - start_date]
  return(chain_pers)
}