R/plot_create_graph.R

In The-Strategy-Unit/723_mh_covid_surge_modelling: Mental Health Covid-19 Surge Modelling

#' Create Graph
#'
#' Generates a network graph of the flows from population groups, to conditions, to treatments, as generated by the
#' model.
#'
#' @param model_output output from \code{run_model()} and \code{get_model_output()}
#' @param groups a vector of population groups to include in the graph, defaults to all groups
#' @param conditions a vector of conditions to include in the graph, defaults to all conditions
#' @param treatments a vector of treatments to include in the graph, defaults to all treatments
#'
#' @return a plotly chart
#'
#' @import rlang
#' @importFrom dplyr %>% filter group_by summarise across bind_rows
#' @importFrom tidyr pivot_longer
#' @importFrom purrr set_names compose map array_tree
#' @importFrom lubridate day
#' @importFrom igraph graph_from_data_frame vertex.attributes
#'   vertex.attributes<- V get.edgelist layout.sugiyama
#' @importFrom plotly plot_ly layout config
create_graph <- function(model_output,
                         groups = unique(model_output$group),
                         conditions = unique(model_output$condition),
                         treatments = unique(model_output$treatment)) {
  df <- model_output %>%
    filter(.data$type == "treatment",
           .data$group %in% groups,
           .data$condition %in% conditions,
           .data$treatment %in% treatments,
           day(.data$date) == 1) %>%
    group_by(.data$group, .data$condition, .data$treatment) %>%
    summarise(across(.data$value, compose(round, sum)), .groups = "drop")

  if (nrow(df) < 1) return(NULL)

  # create a graph of groups to conditions and conditions to the treatment
  # note however, this graph is "reversed", e.g. treatment points to conditions
  # the layout didn't work otherwise.
  g <- bind_rows(
    df %>% group_by(from = .data$condition, to = .data$group),
    df %>% group_by(from = .data$treatment, to = .data$condition)
  ) %>%
    summarise(weight = sum(.data$value), .groups = "drop") %>%
    # remove any lines that after rounding sum to 0
    filter(.data$weight > 0) %>%
    graph_from_data_frame()

  # converts the graph to be a bipartite graph
  vertex.attributes(g)$type <- vertex.attributes(g)$name %in% unique(df$condition)

  # calculate the "weight" of each vertex
  vertex_weights <- df %>%
    pivot_longer(-.data$value, names_to = "type", values_to = "name") %>%
    group_by(.data$type, .data$name) %>%
    summarise(across(.data$value, sum), .groups = "drop")
  # convert to a named list: add in the current treatment as an option also
  vertex_weights <- set_names(vertex_weights$value, vertex_weights$name)

  # set the "weight" attribute of this vertex
  vertex.attributes(g)$weight <- vertex_weights[vertex.attributes(g)$name]

  # extract the vertices
  vs <- V(g)
  # and the edges
  es <- as.data.frame(get.edgelist(g))
  # create a layout for the graph ready to plot
  ly <- layout.sugiyama(g)$layout

  # extract the x- and y-coordinates from the layout
  xs <- ly[, 2]
  ys <- ly[, 1]

  p <- plot_ly(x = ~ xs,
               y = ~ ys,
               size = 1,
               mode = "markers",
               type = "scatter",
               text = paste0("<b>", vs$name, "</b>: ", round(vs$weight)),
               marker = list(
                 opacity = 1,
                 # set the size of each marker to be based on the amount of
                 # people in this group.
                 # take the weight of each vertex and divide by the total size,
                 # so this will convert into the range [0, 1]
                 # take the sqrt of this to make a non-linear increase in size
                 # - this makes smaller dots bigger than they should be
                 # multiply by 500 to make the dots visible (size in pixels?)
                 # ceiling it to turn into an integer
                 size = ceiling(sqrt(unname(vs$weight)) / sqrt(sum(df$value)) * 500)
               ),
               # NHS Blue
               color = I("#005EB8"),
               hoverinfo = "text")

  # iterate over all of the edges and build a shape for each edge: e.g. a line
  # connecting each pair of vertices
  edge_shapes <- map(array_tree(es), function(e) {
    # get the first and second vertex
    v0 <- which(vs$name == e$V1)
    v1 <- which(vs$name == e$V2)

    # create the points
    list(
      type = "line",
      marker = list(color = "#587FC1"),
      line = list(color = "#2c2825", width = 1),
      # ensure th
      layer = "below",
      x0 = xs[v0],
      y0 = ys[v0],
      x1 = xs[v1],
      y1 = ys[v1]
    )
  })

  # render the plot, adding the edges
  layout(p,
         shapes = edge_shapes,
         xaxis = list(visible = FALSE),
         yaxis = list(visible = FALSE)) %>%
    config(displayModeBar = FALSE)
}