R/make_diagram.R

In andreashandel/modeldiagram: Easy Creation of ggplot2 Based Figures of Flow Diagrams

#' Make a ggplot2 model diagram.
#'
#' @description
#' `make_diagram()` generates a **ggplot2** object based on the data frames
#'     made with \code{\link{prepare_diagram}} and, optionally, updated with
#'     \code{\link{update_diagram}}.
#'
#' @param diagram_list A required list of data frames returned from the
#'     \code{\link{prepare_diagram}} function and, optionally, updated with
#'     \code{\link{update_diagram}}. See those functions for details
#'     about this object.
#' @param with_grid A logical indicating whether to return the ggplot
#'     with a grid. Default is FALSE. The grid can be helpful if you
#'     want/need to move items around.
#'
#' @return A ggplot2 object.
#'
#' @details This function uses all the information in the data frames list
#'    generated by \code{\link{prepare_diagram}} and, optionally, updated with
#'    \code{\link{update_diagram}} to make a `ggplot2` object. All location
#'    information and aesthetics are assumed fixed at this point -- no updates
#'    are made within this function. The underlying `ggplot2` code can be
#'    viewed by typing \code{make_diagram} with no parentheses in the R console.
#'
#' @examples
#' mymodel = list(variables = c("S","I","R"),
#'                flows = list(S_flows = c("-b*S*I"),
#'                             I_flows = c("b*S*I","-g*I"),
#'                             R_flows = c("g*I") ) )
#' diagram_list <- prepare_diagram(model_list = mymodel)
#'
#' # make diagram without grid
#' diagram <- make_diagram(diagram_list)
#'
#' # make diagram with grid
#' diagram_with_grid <- make_diagram(diagram_list, with_grid = TRUE)
#'
#' @import ggplot2
#' @export
#'

make_diagram <- function (diagram_list, with_grid = FALSE) {

  # check input data frames for conformity
  test <- check_dataframes(diagram_list)
  if(!is.null(test)) {
    stop(test)
  }

  # unlist the data frames to objects
  variables <- diagram_list$variables
  flows <- diagram_list$flows

  ###
  # make the diagram with ggplot2
  ###
  # Start with an empty ggplot2 canvas. The coord_equal function ensures
  # that the x and y coordinates are displayed in equal proportions to
  # on another (that is, it makes sure that the squares look like squares).
  # All layers are added sequentially onto this blank canvas.
  diagram_plot <- ggplot() +
    coord_equal(clip = "off")


  # LAYER 1: STATE VARIABLES
  # plot the states variable nodes as rectangles

  # The variables data frame is used to create rectangles, with size determined
  # by the xmin, xmax, ymin, and ymax values in the nodes data frame. The
  # outline color of the rectangles is defined by var_outline_color; the
  # inside color (fill) of the rectangles is defined by var_fill_color.
  # The color variables can be a single value or a vector, giving different
  # colors to different rectangles/nodes/state variables. If a vector, the
  # color and fill vectors must have a length that is equal to the number
  # of rows in the nodes data frame (one value for each row).

  # create the nodes/boxes/variables
  # these are just empty rectangles with no text
  for(i in 1:nrow(variables)) {
    diagram_plot <- diagram_plot +  # add new stuff to blank canvas
      geom_rect(
        data = variables[i, ],  # one row of the data frame
        aes(xmin = xmin, xmax = xmax, ymin = ymin, ymax = ymax),  # location information
        color = variables[i, "outline_color"],  # border color
        fill = variables[i, "fill_color"]  # internal, fill color
      )
  }

  # add label text, which goes on top of boxes based on location information
  for(i in 1:nrow(variables)) {
    diagram_plot <- diagram_plot +  # add text to boxes
      geom_text(
        data = variables[i, ],
        aes(x = xlabel, y = ylabel, label = label_text),
        size = variables[i, "label_size"],
        color = variables[i, "label_color"]
      )
  }

  ## add in all the flows
  # start with the lines/arrows
  for(i in 1:nrow(flows)) {
    if(flows[i, "show_arrow"] == TRUE) {
      diagram_plot <- diagram_plot +  # add the lines to the plot with boxes
        geom_curve(  # always use geom_curve, which is straight when cuvature = 1
          data = flows[i, ],
          aes(x = xstart,
              y = ystart,
              xend = xend,
              yend = yend),
          linetype = flows[i, "line_type"],
          arrow = arrow(length = unit(flows[i, "arrow_size"],"cm"), type = "closed"),
          color = flows[i, "line_color"],
          arrow.fill = flows[i, "line_color"],
          lineend = "round",
          linewidth = flows[i, "line_size"],
          curvature = flows[i, "curvature"],
          ncp = 1000  # controls smoothness of curve, larger number = more smooth
        )
    }
  }

  for(i in 1:nrow(flows)) {
    # only plot the label if the arrow is plotted, too
    if(flows[i, "show_label"] == TRUE & flows[i, "show_arrow"] == TRUE) {
      diagram_plot <- diagram_plot +  # now add the flow labels to the canvas
        geom_text(
          data = flows[i, ],
          aes(x = xlabel, y = ylabel, label = label_text),
          size = flows[i, "label_size"],
          color = flows[i, "label_color"])
    }
  }

  # If with_grid == FALSE (default) then void out the theme
  # otherwise keep the grey background with grid
  # the grid can be useful for updating positions of items
  if(with_grid == FALSE) {
    diagram_plot <- diagram_plot +
      theme_void()  # makes an empty plot theme with no axes, grids, or ticks
  } else {
    # The else here may seem silly, but otherwise the returned plot is NULL
    diagram_plot <- diagram_plot  # just returns default ggplot2 theme
  }

  return(diagram_plot)
}