R/plot_functions.R

In rchan26/BayesLogitFusion: Monte Carlo Fusion for Bayesian Logistic Regression

#' @export
plot_fusion_results <- function(plot_rows, 
                                plot_columns, 
                                full_post, 
                                fusion_post, 
                                ylimit = NULL, 
                                dimensions = FALSE,
                                bw = NULL) {
  original_par <- par()
  if (is.null(bw)) {
    bw <- rep("nrd0", ncol(full_post))
  } else if (length(bw)!=ncol(full_post)) {
    stop('plot_fusion_results: bw must be a vector of length ncol(full_post)')
  }
  
  if (dimensions) {
    # first plot each of the dimensions 
    par(mfrow=c(1, 2))
    for (i in 1:ncol(full_post)) {
      plot(density(full_post[,i], bw = bw[i]), xlab = paste(expression(beta),i-1), main = 'full')
      abline(v=mean(full_post[,i]), lty = 4)
      plot(density(fusion_post[,i], bw = bw[i]), xlab = paste(expression(beta),i-1), col = 'red', lty = 2, 'fusion')
      abline(v=mean(full_post[,i]), lty = 4)
    }
  }
  
  par(mfrow=c(plot_rows, plot_columns))
  for (i in 1:ncol(full_post)) {
    if (is.null(ylimit)) {
      plot(density(full_post[,i], bw = bw[i]), xlab = paste(expression(beta),i-1), main = '')
      lines(density(fusion_post[,i], bw = bw[i]), col = 'red', lty = 2)
    } else {
      if (length(ylimit)!=ncol(full_post)) {
        stop('plot_fusion_results: ylimit must be a vector of length ncol(full_post)')
      }
      plot(density(full_post[,i], bw = bw[i]), xlab = paste(expression(beta),i-1), main = '', ylim = c(0, ylimit[i]))
      lines(density(fusion_post[,i], bw = bw[i]), col = 'red', lty = 2)
    }
  }
  
  # reset plots
  par(original_par)
}

#' @export
plot_fusion_matrix <- function(full_post, 
                               fusion_post, 
                               common_limit, 
                               title = NULL,
                               bw = NULL) {
  original_par <- par()
  if (ncol(full_post)!=ncol(fusion_post)) {
    stop('plot_fusion_matrix: full_post and fusion_post must be matrices of same dimension')
  }
  
  # create dimensions of plot
  dimensions <- ncol(full_post)
  # if title is provided, we need to have more space on the top to accommodate
  if (!is.null(title)) {
    par(mfrow=c(dimensions, dimensions), 
        mai = c(0.25, 0.25, 0.25, 0.25),
        oma = c(0.75, 1, 2.5, 1))
  } else {
    par(mfrow=c(dimensions, dimensions), 
        mai = c(0.25, 0.25, 0.25, 0.25),
        oma = c(0.75, 1, 0.75, 1))
  }
  
  if (is.null(bw)) {
    bw <- rep("nrd0", dimensions)
  } else if (length(bw)!=dimensions) {
    stop('plot_fusion_matrix: bw must be a vector of length ncol(full_post)')
  }
  
  for (i in 1:dimensions) {
    for (j in 1:dimensions) {
      # plot the univariate samples
      if (i==j) {
        plot(density(full_post[,i], bw = bw[i]), xlab = '', ylab = '', main = '')
        lines(density(fusion_post[,i], bw = bw[i]), col = 'red', lty = 2)
      } else if (i < j) {
        # get the limits of x and y that just should just fit in the kdes
        xlim_down <- min(min(full_post[,i]), min(fusion_post[,i]))
        xlim_up <- max(max(full_post[,i]), max(fusion_post[,i]))
        ylim_down <- min(min(full_post[,j]), min(fusion_post[,j]))
        ylim_up <- max(max(full_post[,j]), max(fusion_post[,j]))
        # plot kde for full posterior
        plot(ks::kde(full_post[,c(i,j)]), xlab = '', ylab = '', main = '', 
             xlim = c(xlim_down, xlim_up), ylim = c(ylim_down, ylim_up))
        # plot kde for fusion posterior
        plot(ks::kde(fusion_post[,c(i,j)]), col = 'red', add = T)
      } else {
        # plot kde for full posterior
        plot(ks::kde(full_post[,c(i,j)]), xlab = '', ylab = '', main = '', 
             xlim = common_limit, ylim = common_limit)
        # plot kde for fusion posterior
        plot(ks::kde(fusion_post[,c(i,j)]), col = 'red', add = T)
      }
    }
  }
  
  # plot title if provided
  if (!is.null(title)) {
    mtext(title, outer = TRUE, cex = 1.5)
  }
  
  # reset plots
  par(original_par)
}

#' @export
compare_samples_univariate <- function(plot_rows, 
                                       plot_columns, 
                                       posteriors, 
                                       colours, 
                                       ylimit = NULL, 
                                       bw = NULL) {
  original_par <- par()
  if (length(posteriors)!=length(colours)) {
    stop('compare_samples_univariate: the length of posteriors and colours must be the same')
  } 
  original_par <- par()
  if (is.null(bw)) {
    bw <- rep("nrd0", ncol(posteriors[[1]]))
  } else if (length(bw)!=ncol(posteriors[[1]])) {
    stop('compare_samples_univariate: bw must be a vector of length ncol(posteriors[[1]])')
  }
  
  par(mfrow=c(plot_rows, plot_columns))
  for (i in 1:ncol(posteriors[[1]])) {
    if (is.null(ylimit)) {
      plot(density(posteriors[[1]][,i], bw = bw[i]), xlab = paste(expression(beta),i-1), main = '', col = colours[1])
      for (index in 2:length(posteriors)) {
        lines(density(posteriors[[index]][,i], bw = bw[i]), col = colours[index], lty = 2)
      }
    } else {
      if (length(ylimit)!=ncol(posteriors[[1]])) {
        stop('compare_samples_univariate: ylimit must be a vector of length ncol(posteriors[[1]])')
      }
      plot(density(posteriors[[1]][,i], bw = bw[i]), xlab = paste(expression(beta),i-1), main = '', 
           ylim = c(0, ylimit[i]), col = colours[1])
      for (index in 2:length(posteriors)) {
        lines(density(posteriors[[index]][,i], bw = bw[i]), col = colours[index], lty = 2)
      }
    }
  }
  
  # reset plots
  par(original_par)
}

#' @export
compare_samples_bivariate <- function(posteriors, 
                            colours, 
                            common_limit, 
                            title = NULL,
                            bw = NULL) {
  original_par <- par()
  if (length(posteriors)!=length(colours)) {
    stop('compare_samples_bivariate: the length of posteriors and colours must be the same')
  } 
  # create dimensions of plot
  dimensions <- ncol(posteriors[[1]])
  # if title is provided, we need to have more space on the top to accommodate
  if (!is.null(title)) {
    par(mfrow=c(dimensions, dimensions), 
        mai = c(0.25, 0.25, 0.25, 0.25),
        oma = c(0.75, 1, 2.5, 1))
  } else {
    par(mfrow=c(dimensions, dimensions), 
        mai = c(0.25, 0.25, 0.25, 0.25),
        oma = c(0.75, 1, 0.75, 1))
  }
  
  if (is.null(bw)) {
    bw <- rep("nrd0", dimensions)
  } else if (length(bw)!=dimensions) {
    stop('compare_samples_bivariate: bw must be a vector of length ncol(full_post)')
  }
  
  for (i in 1:dimensions) {
    for (j in 1:dimensions) {
      # plot the univariate samples
      if (i==j) {
        plot(density(posteriors[[1]][,i], bw = bw[i]), xlab = '', ylab = '', main = '', col = colours[1])
        for(index in 2:length(posteriors)) {
          lines(density(posteriors[[index]][,i], bw = bw[i]), col = colours[index], lty = index)
        }
      } else if (i < j) {
        # get the limits of x and y that just should just fit in the kdes
        xlim_down <- min(sapply(posteriors, function(post) min(post[,i])))
        xlim_up <- max(sapply(posteriors, function(post) max(post[,i])))
        ylim_down <- min(sapply(posteriors, function(post) min(post[,j])))
        ylim_up <- max(sapply(posteriors, function(post) max(post[,j])))
        # plot kde for full posterior
        plot(ks::kde(posteriors[[1]][,c(i,j)]), xlab = '', ylab = '', main = '', 
             xlim = c(xlim_down, xlim_up), ylim = c(ylim_down, ylim_up), col = colours[1])
        # plot kde for fusion posterior
        for (index in 2:length(posteriors)) {
          plot(ks::kde(posteriors[[index]][,c(i,j)]), col = colours[index], add = T)
        }
      } else {
        # plot kde for full posterior
        plot(ks::kde(posteriors[[1]][,c(i,j)]), xlab = '', ylab = '', main = '', 
             xlim = common_limit, ylim = common_limit, col = colours[1])
        # plot kde for fusion posterior
        for (index in 2:length(posteriors)) {
          plot(ks::kde(posteriors[[index]][,c(i,j)]), col = colours[index], add = T)
        }
      }
    }
  }
  
  # plot title if provided
  if (!is.null(title)) {
    mtext(title, outer = TRUE, cex = 1.5)
  }
  
  # reset plots
  par(original_par)
}

#' @export
integrated_abs_distance <- function(full_post, fusion_post, bw = NULL) {
  if (ncol(full_post)!=ncol(fusion_post)) {
    stop('integrated_abs_distance: full_post and fusion_post must be matrices of same dimension')
  } 
  
  dimensions <- ncol(full_post)
  if (is.null(bw)) {
    bw <- rep("nrd0", dimensions)
  } else if (length(bw)!=dimensions) {
    stop('integrated_abs_distance: bw must be a vector of length ncol(full_post)')
  }
  
  abs_dist <- rep(0, dimensions)
  int_abs_dist <- rep(0, dimensions)
  bandwidth_chosen <- rep(0, dimensions)
  for (i in 1:dimensions) {
    min_value <- min(c(full_post[,i], fusion_post[,i]))
    max_value <- max(c(full_post[,i], fusion_post[,i]))
    # calculate kde for ith dimension in posterior samnples
    baseline_kde <- density(full_post[,i], 
                            bw = bw[i], 
                            from = min_value, 
                            to = max_value,
                            n = 10000)
    bandwidth_chosen[i] <- baseline_kde$bw
    # calculate kde for ith dimension in fusion samples
    fusion_kde <- density(fusion_post[,i], 
                          bw = bandwidth_chosen[i],
                          from = min_value, 
                          to = max_value,
                          n = 10000)
    
    if (!identical(baseline_kde$x, fusion_kde$x)) {
      stop('integrated_abs_distance: the grid for x values are not the same')
    }
    
    # obtain f(x) value from kdes
    baseline_y <- baseline_kde$y
    fusion_y <- fusion_kde$y
    # calculate differences between baseline_y and fusion_y
    diff <- abs(baseline_y - fusion_y)
    # calculating the total variation
    int_abs_dist[i] <- sfsmisc:::integrate.xy(x = baseline_kde$x, fx = diff)
  }
  
  print('bandwidths chosen:'); print(bandwidth_chosen)
  return(sum(int_abs_dist)/(2*dimensions))
}