R/plot.R

In vincentarelbundock/marginsplot: Plot (conditional) marginal effects using the 'margins' and 'ggplot2' packages

#' @title Plot Marginal Effects Estimates
#' @description An implementation of Stata's \samp{marginsplot} as an S3
#'   generic function
#' @param x An object of class \dQuote{margins}, as returned by
#'   \code{\link{margins}}.
#' @param level A numeric value between 0 and 1 indicating the confidence level
#'   to use when drawing error bars.
#' @param term_map named character vector. Values refer to the original
#'   term/variable names. Names refer to the term names that will appear in
#'   the plot. Terms which are omitted from this vector will be omitted from the
#'   plot.
#' @param \dots Additional arguments passed to
#'   \code{\link[ggplot2]{geom_pointrange}}, such as \code{shape}, \code{size},
#'   \code{colour}, etc.
#' @details This function is invoked for its side effect: a basic dot plot with
#'   error bars displaying marginal effects as generated by
#'   \code{\link{margins}}, in the style of Stata's \samp{marginsplot} command.
#' @return A ggplot2 object
#' @examples
#' \dontrun{
#'   require("datasets")
#'   x <- lm(mpg ~ cyl * hp + wt, data = mtcars)
#'   mar <- margins(x)
#'   plot(mar)
#' }
#' 
#' @seealso \code{\link{margins}}, \code{\link{persp.lm}}
#' @keywords graphics
#' @importFrom ggplot2 ggplot aes_string geom_point geom_pointrange theme_minimal xlab ylab
#' @export
plot.margins <- 
function(x, 
         level = 0.95,
         term_map = NULL,
         ...) {
    
    # extract summary
    summ <- summary(x, level = level, by_factor = TRUE)

    # subset and rename
    if (!is.null(term_map)) {
        summ <- summ[summ$factor %in% term_map,]
        if (!is.null(names(term_map))) {
            summ$factor <- names(term_map)[match(summ$factor, term_map)]
        }
    }

    # are confidence intervals available?
    ci <- !all(is.na(summ$lower)) & !all(is.na(summ$upper))
    if (ci) {
        p <- ggplot(summ, aes_string(x = 'factor', y = 'AME', ymin = 'lower', ymax = 'upper')) +
             geom_pointrange(...)
    } else {
        p <- ggplot(summ, aes_string(x = 'factor', y = 'AME')) +
             geom_point(...)
    }

    # finish
    out <- p +
           ylab("Average Marginal Effect") + 
           xlab("")
           theme_minimal() +

    # output
    return(p) 
}