R/plot_sim.R

In multivar: Penalized Estimation of Multiple-Subject Vector Autoregressive (multi-VAR) Models

#' Plot data arising from multivar_sim.
#'
#' @param x Object. An object returned by multivar_sim.
#' @param plot_type Character. User can specify "common" to plot the common effects matrix, "unique" to plot the unique effects matrix, or "total" to plot the total effects matrix.
#' @param facet_ncol Numeric. Number of columns to use in the "unique" or "total" effects plot.
#' @param datasets Numeric. A vector containing the index of datasets to plot. Default is "all".
#' @param ub Numeric. Upper bound on coefficient values for heatmap index. Default is 1.
#' @param lb Numeric. Lower bound on coefficient values for heatmap index. Default is -1.
#' @keywords var multivar simulate plot
#'
#' @examples
#'
#' k <- 3
#' d <- 5
#' n <- 50
#' prop_fill_com <- .2
#' prop_fill_ind <- .2
#' lb <- 0.1
#' ub <- 0.7
#' sigma <- diag(0.1,d)
#' sim <- multivar_sim(k, d, n, prop_fill_com, prop_fill_ind, lb, ub,sigma)
#' plot_sim(sim, plot_type = "common")
#'
#' @export
plot_sim <- function(x, plot_type = "common", facet_ncol=3, datasets = "all",ub = 1, lb = -1){
  rows <- values <- NULL
  if(plot_type == "common"){
    colnames(x$mat_com) <- rownames(x$mat_com) <- paste0("V",1:nrow(x$mat_com))
    df <- setNames(melt(x$mat_com), c('rows', 'vars', 'values'))
  }
  
  if(plot_type == "unique"){
    if(datasets[1] == "all"){
      df <- as.data.frame(do.call("rbind",lapply(seq_along(x$mat_ind_unique), function(j){
        mat <- x$mat_ind_unique[[j]]
        colnames(mat) <- rownames(mat) <- paste0("V",1:nrow(mat))
        df <- setNames(melt(mat), c('rows', 'vars', 'values'))
        df$Subject <- paste0("Dataset ", j)
        df
      })))
    } else {
      df <- as.data.frame(do.call("rbind",lapply(datasets, function(j){
        mat <- x$mat_ind_unique[[j]]
        colnames(mat) <- rownames(mat) <- paste0("V",1:nrow(mat))
        df <- setNames(melt(mat), c('rows', 'vars', 'values'))
        df$Subject <- paste0("Dataset ", j)
        df
      })))
    }
  }
  
  if(plot_type == "total"){
    if(datasets[1] == "all"){
      df <- as.data.frame(do.call("rbind",lapply(seq_along(x$mat_ind_final), function(j){
        mat <- x$mat_ind_final[[j]]
        colnames(mat) <- rownames(mat) <- paste0("V",1:nrow(mat))
        df <- setNames(melt(mat), c('rows', 'vars', 'values'))
        df$Subject <- paste0("Dataset ", j)
        df
      })))
    } else {
      df <- as.data.frame(do.call("rbind",lapply(datasets, function(j){
        mat <- x$mat_ind_final[[j]]
        colnames(mat) <- rownames(mat) <- paste0("V",1:nrow(mat))
        df <- setNames(melt(mat), c('rows', 'vars', 'values'))
        df$Subject <- paste0("Dataset ", j)
        df
      })))
    }
  }

  df$values[df$values == 0] <- NA

  zf_red <- rgb(255,0,90, maxColorValue=255)
  zf_green <- rgb(90, 168, 0, maxColorValue=255)
  zf_blue <- rgb(0, 152, 233, maxColorValue=255)
  zf_yellow <- rgb(242, 147, 24, maxColorValue=255)
  zf_back <- rgb(51,51,51, maxColorValue=255)
  #zf_fore <- rgb(249,242,215, maxColorValue=255)
  zf_fore <- "white"
  
  text_color <- zf_back
  grid_color <- zf_back
  plot_background <- "white"
  
  limit <- max(abs(c(lb,ub))) * c(-1, 1)
  
  colfunc_low <- colorRampPalette(c(zf_red, zf_fore))
  colfunc_high <- colorRampPalette(c(zf_fore, zf_blue))
  colors_to_use <- c(colfunc_low(6)[1:3],zf_fore,colfunc_high(6)[4:6])
  
  df$rows <- factor(df$rows,levels=rev(paste0("V",1:nrow(x$mat_com))))
 
  gg <- ggplot(df, aes(y=rows, x=vars, fill = values)) # original
  #gg <- gg + geom_tile(color=grid_color, size=.5)
  gg <- gg + geom_tile() 
  gg <- gg + scale_fill_gradientn(colors=colors_to_use,limits=limit,na.value = zf_fore,guide = guide_colorbar(frame.colour = "black", ticks.colour = "black",ticks.linewidth = 1,frame.linewidth = 1))
  gg <- gg + coord_equal()
  gg <- gg + theme(panel.grid.minor=element_blank())
  gg <- gg + theme(panel.grid.major=element_blank())
  gg <- gg + theme(axis.text.x = element_text(angle=45,hjust=1))
  gg <- gg + theme(axis.ticks =element_blank())
  gg <- gg + theme(axis.text.x=element_text(size=10, color=text_color))
  gg <- gg + theme(axis.text.y=element_text(size=10, color=text_color))
  gg <- gg + theme(panel.border=element_blank())
  gg <- gg + theme(plot.title=element_text(hjust=0, color=text_color,face="bold"))
  gg <- gg + theme(strip.text=element_text(hjust=0, color=text_color,size=12,face="bold"))
  gg <- gg + theme(strip.background=element_rect(fill=plot_background, color=plot_background))
  gg <- gg + theme(panel.spacing.x=unit(0.5, "cm"))
  gg <- gg + theme(panel.spacing.y=unit(0.5, "cm"))
  gg <- gg + theme(legend.background=element_rect(fill=plot_background, color=plot_background)) 
  gg <- gg + theme(legend.title=element_text(size=12, color=text_color))
  gg <- gg + theme(legend.title.align=1)
  gg <- gg + theme(legend.text=element_text(size=10, color=text_color))
  gg <- gg + theme(plot.background=element_rect(fill=plot_background,color=plot_background)) 
  gg <- gg + theme(panel.border=element_rect(fill = NA, colour='black',size=1))
  gg <- gg + theme(legend.text.align=1)
  
  gg <- gg + labs(fill='') 
  
  if(plot_type == "common"){
     gg <- gg + labs(x=NULL, y=NULL, title="Common Transition Matrix")
  }
  
  if(plot_type == "unique"){
     gg <- gg + labs(x=NULL, y=NULL, title="Unique Transition Matrices")
     gg <- gg + facet_wrap(~Subject, ncol=facet_ncol)
  }
  
  if(plot_type == "total"){
     gg <- gg + labs(x=NULL, y=NULL, title="Total Transition Matrices")
     gg <- gg + facet_wrap(~Subject, ncol=facet_ncol)
  }
  
  gg
  
  
}