-
Notifications
You must be signed in to change notification settings - Fork 24
/
vip.R
91 lines (89 loc) · 3.1 KB
/
vip.R
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#' Variable Importance Plots
#'
#' Plot variable importance scores for the predictors in a model.
#'
#' @param object A fitted model object (e.g., a \code{"randomForest"} object).
#'
#' @param num_features Integer specifying the number of variable importance
#' scores to plot. Default is \code{10}.
#'
#' @param bar Logical indicating whether or not to produce a barplot. Default
#' is \code{TRUE}. If \code{bar = FALSE}, then a dotchart is displayed instead.
#'
#' @param width Numeric value specifying the width of the bars when
#' \code{bar = TRUE}. Default is \code{0.75}.
#'
#' @param horizontal Logical indicating whether or not to plot the importance
#' scores on the x-axis (\code{TRUE}). Default is \code{TRUE}.
#'
#' @param alpha Numeric value between 0 and 1 giving the transparency of the
#' bars (\code{bar = TRUE}) or points (\code{bar = FALSE}).
#'
#' @param color Character string specifying the color to use for the borders of
#' the bars. Could also be a function, such as
#' \code{\link[grDevices]{heat.colors}}. Default is \code{"grey35"}.
#'
#' @param fill Character string specifying the color to use to fill the bars.
#' Could also be a function, such as \code{\link[grDevices]{heat.colors}}.
#' Default is \code{"grey35"}.
#'
#' @param size Numeric value indicating the size to use for the points whenever
#' \code{bar = FALSE}. Default is \code{1}.
#'
#' @param shape Numeric value indicating the shape to use for the points
#' whenever \code{bar = FALSE}. Default is \code{1}.
#'
#' @param ... Additional optional arguments to be passed onto \code{\link{vi}}.
#'
#' @rdname vip
#'
#' @export
#'
#' @examples
#' #
#' # A projection pursuit regression example
#' #
#'
#' # Load the sample data
#' data(mtcars)
#'
#' # Fit a projection pursuit regression model
#' mtcars.ppr <- ppr(mpg ~ ., data = mtcars, nterms = 1)
#'
#' # Construct variable importance plot
#' vip(mtcars.ppr, method = "ice")
vip <- function(object, ...) {
UseMethod("vip")
}
#' @rdname vip
#'
#' @export
vip.default <- function(
object, num_features = 10L, bar = TRUE, width = 0.75, horizontal = TRUE,
alpha = 1, color = "grey35", fill = "grey35", size = 1, shape = 19, ...
) {
imp <- vi(object, ...) # variable importance scores
vi_type <- attr(imp, which = "type") # subsetting removes this attribute!
num_features <- as.integer(num_features)[1L] # make sure num_features is a single integer
if (num_features > nrow(imp) || num_features < 1L) {
num_features <- nrow(imp)
}
imp <- imp[seq_len(num_features), ] # only retain num_features variable importance scores
x.string <- "reorder(Variable, Importance)"
p <- ggplot2::ggplot(imp, ggplot2::aes_string(x = x.string, y = "Importance"))
p <- if (bar) {
p + ggplot2::geom_col(
width = width, color = color, fill = fill, alpha = alpha
)
} else {
p + ggplot2::geom_point(
color = color, alpha = alpha, size = size, shape = shape
)
}
p <- p + ggplot2::theme(legend.position = "none")
p <- p + ggplot2::xlab("") # no need for x-axis label
if (horizontal) {
p <- p + ggplot2::coord_flip()
}
p + ggplot2::ylab(paste0("Importance (", vi_type, ")"))
}